cp_library/ds/tree/seg/segtree2_cls.py

• View this file on GitHub
• Last update: 2025-08-24 20:27:05+09:00

Depends on

• argsort (cp_library/alg/iter/arg/argsort_fn.py)
• cp_library/alg/iter/sort/isort_parallel_fn.py
• cp_library/bit/pack/packer_cls.py
• cp_library/ds/list/list2_cls.py
• cp_library/ds/tree/seg/segtree_cls.py

Required by

• perf/segtree2.py

Verified with

• test/unittests/ds/tree/seg/segtree2_cls_test.py

Code

import cp_library.__header__
from cp_library.misc.typing import _T
import cp_library.ds.__header__
from cp_library.ds.list.list2_cls import list2
import cp_library.ds.tree.__header__
import cp_library.ds.tree.seg.__header__
from cp_library.ds.tree.seg.segtree_cls import SegTree

class SegTree2(SegTree[_T]):
    _lst = list2

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

_S = TypeVar('S'); _T = TypeVar('T'); _U = TypeVar('U'); _T1 = TypeVar('T1'); _T2 = TypeVar('T2'); _T3 = TypeVar('T3'); _T4 = TypeVar('T4'); _T5 = TypeVar('T5'); _T6 = TypeVar('T6')

from typing import Generic




def argsort(A: list[int], reverse=False):
    P = Packer(len(I := list(A))-1); P.ienumerate(I, reverse); I.sort(); P.iindices(I)
    return I



class Packer:
    __slots__ = 's', 'm'
    def __init__(P, mx: int): P.s = mx.bit_length(); P.m = (1 << P.s) - 1
    def enc(P, a: int, b: int): return a << P.s | b
    def dec(P, x: int) -> tuple[int, int]: return x >> P.s, x & P.m
    def enumerate(P, A, reverse=False): P.ienumerate(A:=list(A), reverse); return A
    def ienumerate(P, A, reverse=False):
        if reverse:
            for i,a in enumerate(A): A[i] = P.enc(-a, i)
        else:
            for i,a in enumerate(A): A[i] = P.enc(a, i)
    def indices(P, A: list[int]): P.iindices(A:=list(A)); return A
    def iindices(P, A):
        for i,a in enumerate(A): A[i] = P.m&a


def isort_parallel(*L: list, reverse=False):
    inv, order = [0]*len(L[0]), argsort(L[0], reverse=reverse)
    for i, j in enumerate(order): inv[j] = i
    for i, j in enumerate(order):
        for A in L: A[i], A[j] = A[j], A[i]
        order[inv[i]], inv[j] = j, inv[i]
    return L


class list2(Generic[_T1, _T2]):
    __slots__ = 'A1', 'A2'
    def __init__(lst, A1: list[_T1], A2: list[_T2]): lst.A1, lst.A2 = A1, A2
    def __len__(lst): return len(lst.A1)
    def __getitem__(lst, i: int): return lst.A1[i], lst.A2[i]
    def __setitem__(lst, i: int, v: tuple[_T1, _T2]): lst.A1[i], lst.A2[i] = v
    def __contains__(lst, v: tuple[_T1, _T2]): raise NotImplementedError
    def index(lst, v: tuple[_T1, _T2]): raise NotImplementedError
    def reverse(lst): lst.A1.reverse(); lst.A2.reverse()
    def sort(lst, reverse=False): isort_parallel(lst.A1, lst.A2, reverse=reverse)
    def pop(lst): return lst.A1.pop(), lst.A2.pop()
    def append(lst, v: tuple[_T1, _T2]): v1, v2 = v; lst.A1.append(v1); lst.A2.append(v2)
    def add(lst, i: int, v: tuple[_T1, _T2]): lst.A1[i] += v[0]; lst.A2[i] += v[1]


from typing import Callable, Generic, Union

class SegTree(Generic[_T]):
    _lst = list
    
    def __init__(seg, op: Callable[[_T, _T], _T], e: _T, v: Union[int, list[_T]]) -> None:
        if isinstance(v, int): n = v; v = None
        else: n = len(v)
        seg.op, seg.e, seg.n = op, e, n
        seg.log, seg.sz = (log := (n-1).bit_length()+1), (sz := 1 << log)
        if seg._lst is list: seg.d = [e]*(sz<<1)
        else: seg.d = seg._lst(*([e_]*(sz<<1) for e_ in e))
        if v: seg._build(v)

    def _build(seg, v):
        for i in range(seg.n): seg.d[seg.sz + i] = v[i]
        for i in range(seg.sz-1,0,-1): seg._merge(i, i<<1, i<<1|1)
    
    def _merge(seg, i, j, k): seg.d[i] = seg.op(seg.d[j], seg.d[k])

    def set(seg, p: int, x: _T) -> None:
        p += seg.sz
        seg.d[p] = x
        for _ in range(seg.log):
            p = p^(p&1)
            seg._merge(p>>1, p, p|1)
            p >>= 1
    __setitem__ = set

    def get(seg, p: int) -> _T: return seg.d[p+seg.sz]
    __getitem__ = get

    def prod(seg, l: int, r: int) -> _T:
        sml = smr = seg.e
        l, r = l+seg.sz, r+seg.sz
        while l < r:
            if l&1: sml, l = seg.op(sml, seg.d[l]), l+1
            if r&1: smr = seg.op(seg.d[r:=r-1], smr)
            l, r = l >> 1, r >> 1
        return seg.op(sml, smr)

    def all_prod(seg) -> _T: return seg.d[1]

    def max_right(seg, l: int, f: Callable[[_T], bool]) -> int:
        assert 0 <= l <= seg.n
        assert f(seg.e)
        if l == seg.n: return seg.n
        l, op, d, sm = l+(sz := seg.sz), seg.op, seg.d, seg.e
        while True:
            while l&1 == 0: l >>= 1
            if not f(op(sm, d[l])):
                while l < sz:
                    if f(op(sm, d[l:=l<<1])): sm, l = op(sm, d[l]), l+1
                return l - sz
            sm, l = op(sm, d[l]), l+1
            if l&-l == l: return seg.n

    def min_left(seg, r: int, f: Callable[[_T], bool]) -> int:
        assert 0 <= r <= seg.n
        assert f(seg.e)
        if r == 0: return 0
        r, op, d, sm = r+(sz := seg.sz), seg.op, seg.d, seg.e
        while True:
            r -= 1
            while r > 1 and r & 1: r >>= 1
            if not f(op(d[r], sm)):
                while r < sz:
                    if f(op(d[r:=r<<1|1], sm)): sm, r = op(d[r], sm), r-1
                return r + 1 - sz
            sm = op(d[r], sm)
            if (r & -r) == r: return 0

class SegTree2(SegTree[_T]):
    _lst = list2

Back to top page