cp_library/math/invcnt_fn.py

• View this file on GitHub
• Last update: 2025-03-30 20:17:47+09:00

Depends on

• cp_library/bit/pack_sm_fn.py
• cp_library/ds/tree/bit/bit_cls.py

Verified with

• test/atcoder/arc/arc136_b_inversion_cnt_fn.test.py

Code

import cp_library.math.__header__
from cp_library.ds.tree.bit.bit_cls import BIT
from cp_library.bit.pack_sm_fn import pack_sm, pack_indices

def invcnt(A: list[int]):
    s, m = pack_sm(N := len(A))
    bit, cnt, I = BIT(N), 0, pack_indices(A, s)
    I.sort(reverse=True)
    for i in I:
        cnt += bit.sum(i&m)
        bit.add(i&m, 1)
    return cnt

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from typing import Union


'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
            ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓            
            ┃                                    7 ┃            
            ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━┛            
            ┏━━━━━━━━━━━━━━━━━━┓                 │              
            ┃                3 ┃◄────────────────┤              
            ┗━━━━━━━━━━━━━━━━┯━┛                 │              
            ┏━━━━━━━━┓       │  ┏━━━━━━━━┓       │              
            ┃      1 ┃◄──────┤  ┃      5 ┃◄──────┤              
            ┗━━━━━━┯━┛       │  ┗━━━━━━┯━┛       │              
            ┏━━━┓  │  ┏━━━┓  │  ┏━━━┓  │  ┏━━━┓  │              
            ┃ 0 ┃◄─┤  ┃ 2 ┃◄─┤  ┃ 4 ┃◄─┤  ┃ 6 ┃◄─┤              
            ┗━┯━┛  │  ┗━┯━┛  │  ┗━┯━┛  │  ┗━┯━┛  │              
              │    │    │    │    │    │    │    │              
              ▼    ▼    ▼    ▼    ▼    ▼    ▼    ▼              
            ┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓┏━━━┓            
            ┃ 0 ┃┃ 1 ┃┃ 2 ┃┃ 3 ┃┃ 4 ┃┃ 5 ┃┃ 6 ┃┃ 7 ┃            
            ┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛┗━━━┛            
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
           Data Structure - Tree - Binary Index Tree            
'''

class BIT:
    def __init__(bit, v: Union[int, list[int]]):
        if isinstance(v, int): bit.d, bit.n = [0]*v, v
        else: bit.build(v)
        bit.lb = 1<<(bit.n.bit_length()-1)

    def build(bit, data):
        bit.d, bit.n = data, len(data)
        for i in range(bit.n):
            if (r := i|i+1) < bit.n: bit.d[r] += bit.d[i]

    def add(bit, i, x):
        while i < bit.n:
            bit.d[i] += x
            i |= i+1

    def sum(bit, n: int) -> int:
        s = 0
        while n: s, n = s+bit.d[n-1], n&n-1
        return s

    def range_sum(bit, l, r):
        s = 0
        while r: s, r = s+bit.d[r-1], r&r-1
        while l: s, l = s-bit.d[l-1], l&l-1
        return s

    def __len__(bit) -> int:
        return bit.n
    
    def __getitem__(bit, i: int) -> int:
        s, l = bit.d[i], i&(i+1)
        while l != i: s, i = s-bit.d[i-1], i-(i&-i)
        return s
    get = __getitem__
    
    def __setitem__(bit, i: int, x: int) -> None:
        bit.add(i, x-bit[i])
    set = __setitem__

    def prelist(bit) -> list[int]:
        pre = [0]+bit.d
        for i in range(bit.n+1): pre[i] += pre[i&i-1]
        return pre

    def bisect_left(bit, v) -> int:
        return bit.bisect_right(v-1) if v>0 else 0
    
    def bisect_right(bit, v) -> int:
        i = s = 0; ni = m = bit.lb
        while m:
            if ni <= bit.n and (ns:=s+bit.d[ni-1]) <= v: s, i = ns, ni
            ni = (m:=m>>1)|i
        return i


def pack_sm(N: int):
    s = N.bit_length()
    return s, (1<<s)-1

def pack_enc(a: int, b: int, s: int):
    return a << s | b
    
def pack_dec(ab: int, s: int, m: int):
    return ab >> s, ab & m

def pack_indices(A, s):
    return [a << s | i for i,a in enumerate(A)]

def invcnt(A: list[int]):
    s, m = pack_sm(N := len(A))
    bit, cnt, I = BIT(N), 0, pack_indices(A, s)
    I.sort(reverse=True)
    for i in I:
        cnt += bit.sum(i&m)
        bit.add(i&m, 1)
    return cnt

Back to top page