cp_library/ds/tree/bitset_tree_cls.py

• View this file on GitHub
• Last update: 2025-05-06 22:58:43+09:00

Depends on

• cp_library/alg/dp/min2_fn.py
• cp_library/ds/array_init_fn.py

Verified with

• test/library-checker/data-structure/predecessor_problem.test.py

Code

from cp_library.ds.array_init_fn import u32f
import cp_library.__header__
from cp_library.alg.dp.min2_fn import min2
import cp_library.ds.__header__
import cp_library.ds.tree.__header__

class BitsetTree:
    def __init__(B, S: str):
        B.N = N = len(S)
        L = u32f((N+31)>>5)
        for k,c in enumerate(S):
            k,r = k>>5,k&31
            if c=='1': L[k]|=1<<r
        B.levels = [L]
        while (N := len(L := B.levels[-1])) > 1:
            A = u32f((N+31)>>5)
            for i in range(0,N,32):
                a=j=0;r=min2(i+32,N)
                while i+j<r:a|=(0<L[i+j])<<j;j+=1
                A[i>>5]=a
            B.levels.append(A)
        B.depth = len(B.levels)

    def set0(B, k):
        if B.levels[0][k>>5]>>(k&31)&1: 
            l = -1
            while (l:=l+1) < B.depth:
                B.levels[l][k>>5]&=~(1<<(k&31));k>>=5
                if B.levels[l][k]: break

    def set1(B, k):
        if not B.levels[0][k>>5]>>(k&31)&1: 
            l = -1
            while (l:=l+1) < B.depth: B.levels[l][k>>5]|=1<<(k&31);k>>=5

    def __setitem__(B, k: int, v: int):
        if v: B.set1(k)
        else: B.set0(k)

    def __getitem__(B, k: int):
        b,r=k>>5,k&31
        return B.levels[0][b]>>r&1
    
    def ge(B, k: int):
        if not B.levels[-1][0]: return -1
        l = -1
        while True:
            k,r=k>>5,k&31
            if m:=(B.levels[l:=l+1][k]>>r)<<r: break
            if (k:=k+1) >= len(B.levels[l]): return -1
        while l:m=B.levels[l:=l-1][k:=k<<5|(m&-m).bit_length()-1]
        return k<<5|(m&-m).bit_length()-1
    
    def le(B, k: int):
        if not B.levels[-1][0]: return -1
        l = -1
        while True:
            k,r=k>>5,k&31
            if m:=B.levels[l:=l+1][k]&((1<<(r+1))-1): break
            if (k:=k-1)<0: return -1
        while l:m=B.levels[l:=l-1][k:=k<<5|m.bit_length()-1]
        return k<<5|m.bit_length()-1

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

def i8f(N: int, elm: int = 0):      return array('b', (elm,))*N  # signed char
def u8f(N: int, elm: int = 0):      return array('B', (elm,))*N  # unsigned char
def i16f(N: int, elm: int = 0):     return array('h', (elm,))*N  # signed short
def u16f(N: int, elm: int = 0):     return array('H', (elm,))*N  # unsigned short
def i32f(N: int, elm: int = 0):     return array('i', (elm,))*N  # signed int
def u32f(N: int, elm: int = 0):     return array('I', (elm,))*N  # unsigned int
def i64f(N: int, elm: int = 0):     return array('q', (elm,))*N  # signed long long
def u64f(N: int, elm: int = 0):     return array('Q', (elm,))*N  # unsigned long long
def f32f(N: int, elm: float = 0.0): return array('f', (elm,))*N  # float
def f64f(N: int, elm: float = 0.0): return array('d', (elm,))*N  # double

def i8a(init = None):  return array('b') if init is None else array('b', init)  # signed char
def u8a(init = None):  return array('B') if init is None else array('B', init)  # unsigned char
def i16a(init = None): return array('h') if init is None else array('h', init)  # signed short
def u16a(init = None): return array('H') if init is None else array('H', init)  # unsigned short
def i32a(init = None): return array('i') if init is None else array('i', init)  # signed int
def u32a(init = None): return array('I') if init is None else array('I', init)  # unsigned int
def i64a(init = None): return array('q') if init is None else array('q', init)  # signed long long
def u64a(init = None): return array('Q') if init is None else array('Q', init)  # unsigned long long
def f32a(init = None): return array('f') if init is None else array('f', init)  # float
def f64a(init = None): return array('d') if init is None else array('d', init)  # double

i8_max = (1 << 7)-1
u8_max = (1 << 8)-1
i16_max = (1 << 15)-1
u16_max = (1 << 16)-1
i32_max = (1 << 31)-1
u32_max = (1 << 32)-1
i64_max = (1 << 63)-1
u64_max = (1 << 64)-1


def min2(a, b):
    return a if a < b else b


class BitsetTree:
    def __init__(B, S: str):
        B.N = N = len(S)
        L = u32f((N+31)>>5)
        for k,c in enumerate(S):
            k,r = k>>5,k&31
            if c=='1': L[k]|=1<<r
        B.levels = [L]
        while (N := len(L := B.levels[-1])) > 1:
            A = u32f((N+31)>>5)
            for i in range(0,N,32):
                a=j=0;r=min2(i+32,N)
                while i+j<r:a|=(0<L[i+j])<<j;j+=1
                A[i>>5]=a
            B.levels.append(A)
        B.depth = len(B.levels)

    def set0(B, k):
        if B.levels[0][k>>5]>>(k&31)&1: 
            l = -1
            while (l:=l+1) < B.depth:
                B.levels[l][k>>5]&=~(1<<(k&31));k>>=5
                if B.levels[l][k]: break

    def set1(B, k):
        if not B.levels[0][k>>5]>>(k&31)&1: 
            l = -1
            while (l:=l+1) < B.depth: B.levels[l][k>>5]|=1<<(k&31);k>>=5

    def __setitem__(B, k: int, v: int):
        if v: B.set1(k)
        else: B.set0(k)

    def __getitem__(B, k: int):
        b,r=k>>5,k&31
        return B.levels[0][b]>>r&1
    
    def ge(B, k: int):
        if not B.levels[-1][0]: return -1
        l = -1
        while True:
            k,r=k>>5,k&31
            if m:=(B.levels[l:=l+1][k]>>r)<<r: break
            if (k:=k+1) >= len(B.levels[l]): return -1
        while l:m=B.levels[l:=l-1][k:=k<<5|(m&-m).bit_length()-1]
        return k<<5|(m&-m).bit_length()-1
    
    def le(B, k: int):
        if not B.levels[-1][0]: return -1
        l = -1
        while True:
            k,r=k>>5,k&31
            if m:=B.levels[l:=l+1][k]&((1<<(r+1))-1): break
            if (k:=k-1)<0: return -1
        while l:m=B.levels[l:=l-1][k:=k<<5|m.bit_length()-1]
        return k<<5|m.bit_length()-1

Back to top page