cp_library/ds/tree/bst/treap_monoid_cls.py

• View this file on GitHub
• Last update: 2025-06-20 03:24:59+09:00

Depends on

• cp_library/bit/masks/i64_max_cnst.py
• cp_library/ds/reserve_fn.py
• cp_library/ds/tree/bst/bst_cls.py
• cp_library/ds/tree/bst/bst_updates_cls.py
• cp_library/ds/tree/bst/cartesian_tree_cls.py
• cp_library/ds/tree/bst/treap_cls.py

Required by

• cp_library/ds/tree/bst/treap_monoid_reversible_cls.py

Verified with

• test/library-checker/data-structure/point_set_range_composite_large_array_treap.test.py
• test/library-checker/data-structure/range_reverse_range_sum.test.py
• test/unittests/ds/tree/bst/treap_monoid_cls_test.py

Code

import cp_library.__header__
import cp_library.ds.__header__
from cp_library.ds.reserve_fn import reserve
import cp_library.ds.tree.__header__
import cp_library.ds.tree.bst.__header__
from cp_library.ds.tree.bst.bst_updates_cls import BSTUpdates
from cp_library.ds.tree.bst.treap_cls import Treap

class TreapMonoid(Treap, BSTUpdates):
    __slots__='op'
    K,V,A,P,sub,st=[-1],[-1],[-1],[42],[-1,-1],[]
    def __init__(T,op,e=-1):T.op=op;super().__init__(e)
    def _nt(T):return T.__class__(T.op,T.e)
    def _nr(T):T.A.append(T.e);return super()._nr()
    def _nn(T,k,v):T.A.append(v);return super()._nn(k, v)
    def prod(T,l,r):
        # find common ancestor
        a=T.sub[T.r<<1]
        while~a and not l<=T.K[a]<r:T._p(a);a=T.sub[a<<1|(T.K[a]<l)]
        if a<0:return T.e
        # left subtreap
        ac,i=T.V[a],T.sub[a<<1]
        while~i:
            T._p(i)
            if not(b:=T.K[i]<l):
                if~(j:=T.sub[i<<1|1]):ac=T.op(T.A[j],ac)
                ac=T.op(T.V[i],ac)
            i=T.sub[i<<1|b]
        # right subtreap
        i=T.sub[a<<1|1]
        while~i:
            T._p(i)
            if b:=T.K[i]<r:
                if~(j:=T.sub[i<<1]):ac=T.op(ac,T.A[j])
                ac=T.op(ac,T.V[i])
            i=T.sub[i<<1|b]
        return ac
    def all_prod(T):return T.A[T.r]
    def __getitem__(T,k):
        if isinstance(k,int):return T.get(k)
        elif isinstance(k,slice):return T.prod(k.start,k.stop)
    @classmethod
    def reserve(cls,sz):super(TreapMonoid,cls).reserve(sz);reserve(cls.A,sz+1)
    def _u(T,i):
        T.A[i]=T.V[i]
        if~(l:=T.sub[i<<1]):T.A[i]=T.op(T.A[l],T.A[i])
        if~(r:=T.sub[i<<1|1]):T.A[i]=T.op(T.A[i],T.A[r])
    def _v(T,i=None):
        if i is None:
            assert T.all_prod() == (ac := T._v(i) if ~(i := T.sub[T.r<<1]) else T.e)
            return ac
        T._p(i);ac = T.V[i]
        if ~(l:=T.sub[i<<1]):
            assert T.P[i] <= T.P[l]
            assert T.K[l] <= T.K[i]
            ac = T.op(T._v(l), ac)
        if ~(r:=T.sub[i<<1|1]):
            assert T.P[i] <= T.P[r]
            assert T.K[i] <= T.K[r]
            ac = T.op(ac, T._v(r))
        assert T.A[i] == ac
        return ac

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


def reserve(A: list, est_len: int) -> None: ...
try:
    from __pypy__ import resizelist_hint
except:
    def resizelist_hint(A: list, est_len: int):
        pass
reserve = resizelist_hint




i64_max = (1<<63)-1

class BST:
    __slots__ = 'r'
    K,sub,st=[-1],[-1,-1],[]
    def __init__(T):T.r=T._nr()
    def _nt(T):return T.__class__()
    def _nr(T):r=len(T.K);T.K.append(i64_max);T.sub.append(-1);T.sub.append(-1);return r
    def _nn(T,k):n=len(T.K);T.K.append(k);T.sub.append(-1);T.sub.append(-1);return n
    def insert(T,k):T._i(T.r<<1,k,n:=T._nn(k));T._r();return n
    def get(T,k):
        if~(i:=T._f(T.r<<1,k)):return i
        raise KeyError
    def pop(T,k):
        if ~(i:=T._t(T.r<<1,k)):T._d(i,T.st[-1]);T._r();return i
        else:T.st.clear();raise KeyError
    def __delitem__(T,k):
        if~(i:=T._t(T.r<<1,k)):T._d(i,T.st[-1]);T._r()
        else:T.st.clear();raise KeyError
    def __contains__(T,k):return 0<=T._f(T.r<<1,k)
    def _f(T,s,k):
        i = T.sub[s]
        while~i and T.K[i]!=k:T._p(i);i=T.sub[i<<1|(T.K[i]<k)]
        return i
    def _t(T,s,k):
        T.st.append(s)
        while~(i:=T.sub[s])and T.K[i]!=k:T._p(i);T.st.append(s:=i<<1|(T.K[i]<k))
        return i
    def _i(T,s,k,n):
        T.st.append(s)
        while ~T.sub[s]:T._p(i:=T.sub[s]);T.st.append(s:=i<<1|(T.K[i]<k))
        i,T.sub[s]=T.sub[s],n
    def _d(T,i,s): raise NotImplemented
    def _r(T):T.st.clear()
    def _p(T,i): pass
    @classmethod
    def reserve(cls,sz):sz+=1;reserve(cls.K,sz);reserve(cls.sub,sz<<1);reserve(cls.st,sz.bit_length()<<1)
    def _node_str(T, i): return f"{T.K[i]}"
    def __str__(T):
        def rec(i, pre="", is_right=False):
            if i == -1: return ""
            ret = "";T._p(i)
            if ~(r:=T.sub[i<<1|1]):ret+=rec(r,pre+("   "if is_right else"│  "),True)
            ret+=pre+("┌─ "if is_right else"└─ ")+T._node_str(i)+"\n"
            if ~(l:=T.sub[i<<1]):ret+=rec(l,pre+("   "if not is_right else"│  "),False)
            return ret
        return rec(T.sub[T.r<<1]).rstrip()

class BSTUpdates(BST):
    def _u(T,i): pass
    def _r(T):
        while T.st:T._u(T.st.pop()>>1)

class CartesianTree(BST):
    K,P,sub,st=[-1],[42],[-1,-1],[]
    def _nr(T):T.P.append(-1);return super()._nr()
    def _nn(T,k,p=-1):T.P.append(p);return super()._nn(k)
    def get(T,k):return T.P[BST.get(T,k)]
    def pop(T,k):return T.P[BST.pop(T,k)]
    def split(T,k):S=T._nt();T._sp(T.sub[T.r<<1],k,S.r<<1,T.r<<1);T._r();return S,T
    def insert(T,k,p):T._i(T.r<<1,k,n:=T._nn(k,p));T._r();return n
    def __getitem__(T,k):return T.get(k)
    def _i(T,s,k,n):
        T.st.append(s)
        while~T.sub[s]and T.P[i:=T.sub[s]]<T.P[n]:T._p(i);T.st.append(s:=i<<1|(T.K[i]<k))
        i,T.sub[s]=T.sub[s],n
        if~i:T._sp(i,k,n<<1,n<<1|1)
    def _sp(T,i,k,l,r):
        T.st.append(l)
        if 1<l^r:T.st.append(r)
        while~i:
            T._p(i)
            if T.K[i]<k:T.sub[l]=i;i=T.sub[l:=i<<1|1];T.st.append(l)
            else:T.sub[r]=i;i=T.sub[r:=i<<1];T.st.append(r)
        T.sub[l]=T.sub[r]=-1
    def _m(T,s,l,r):
        T.st.append(s)
        while~l and~r:
            if T.P[l]<T.P[r]:T._p(l);T.sub[s]=l;l=T.sub[s:=l<<1|1]
            else:T._p(r);T.sub[s]=r;r=T.sub[s:=r<<1]
            T.st.append(s)
        T.sub[s]=l if~l else r
    def _d(T,i,s):T._p(i);T._m(s,T.sub[i<<1],T.sub[i<<1|1])
    @classmethod
    def reserve(cls,sz):super(CartesianTree,cls).reserve(sz);reserve(cls.P,sz+1)

class Treap(CartesianTree):
    __slots__='e'
    K,V,P,sub,st=[-1],[-1],[42],[-1,-1],[]
    def __init__(T,e=-1):T.e=e;super().__init__()
    def _nt(T):return T.__class__(T.e)
    def _nr(T):T.V.append(T.e);return super()._nr()
    def _nn(T,k,v):T.V.append(v);return super()._nn(k,(T.P[-1]*1103515245+12345)&0x7fffffff)
    def insert(T,k,v):return super().insert(k,v)
    def get(T,k):return T.V[BST.get(T,k)]
    def pop(T,k):return T.V[BST.pop(T,k)]
    def set(T,k,v):T._s(T.r<<1,k,v);T._r()
    def __setitem__(T,k,v):T.set(k,v)
    def _s(T,s,k,v):
        if ~(i:=T._t(s,k)):T.V[i]=v;T.st.append(i<<1)
        else:
            n=T._nn(k,v)
            while T.P[n]<T.P[i:=T.st[-1]>>1]:T._p(T.st.pop())
            T._p(i)
            i,T.sub[s]=T.sub[s:=i<<1|(i!=T.r and T.K[i]<k)],n
            if~i:T._sp(i,k,n<<1,n<<1|1)
    def _node_str(T, i): return f"{T.K[i]}:{T.V[i]}"
    @classmethod
    def reserve(cls,hint):super(Treap,cls).reserve(hint);reserve(cls.V,hint+1)

class TreapMonoid(Treap, BSTUpdates):
    __slots__='op'
    K,V,A,P,sub,st=[-1],[-1],[-1],[42],[-1,-1],[]
    def __init__(T,op,e=-1):T.op=op;super().__init__(e)
    def _nt(T):return T.__class__(T.op,T.e)
    def _nr(T):T.A.append(T.e);return super()._nr()
    def _nn(T,k,v):T.A.append(v);return super()._nn(k, v)
    def prod(T,l,r):
        # find common ancestor
        a=T.sub[T.r<<1]
        while~a and not l<=T.K[a]<r:T._p(a);a=T.sub[a<<1|(T.K[a]<l)]
        if a<0:return T.e
        # left subtreap
        ac,i=T.V[a],T.sub[a<<1]
        while~i:
            T._p(i)
            if not(b:=T.K[i]<l):
                if~(j:=T.sub[i<<1|1]):ac=T.op(T.A[j],ac)
                ac=T.op(T.V[i],ac)
            i=T.sub[i<<1|b]
        # right subtreap
        i=T.sub[a<<1|1]
        while~i:
            T._p(i)
            if b:=T.K[i]<r:
                if~(j:=T.sub[i<<1]):ac=T.op(ac,T.A[j])
                ac=T.op(ac,T.V[i])
            i=T.sub[i<<1|b]
        return ac
    def all_prod(T):return T.A[T.r]
    def __getitem__(T,k):
        if isinstance(k,int):return T.get(k)
        elif isinstance(k,slice):return T.prod(k.start,k.stop)
    @classmethod
    def reserve(cls,sz):super(TreapMonoid,cls).reserve(sz);reserve(cls.A,sz+1)
    def _u(T,i):
        T.A[i]=T.V[i]
        if~(l:=T.sub[i<<1]):T.A[i]=T.op(T.A[l],T.A[i])
        if~(r:=T.sub[i<<1|1]):T.A[i]=T.op(T.A[i],T.A[r])
    def _v(T,i=None):
        if i is None:
            assert T.all_prod() == (ac := T._v(i) if ~(i := T.sub[T.r<<1]) else T.e)
            return ac
        T._p(i);ac = T.V[i]
        if ~(l:=T.sub[i<<1]):
            assert T.P[i] <= T.P[l]
            assert T.K[l] <= T.K[i]
            ac = T.op(T._v(l), ac)
        if ~(r:=T.sub[i<<1|1]):
            assert T.P[i] <= T.P[r]
            assert T.K[i] <= T.K[r]
            ac = T.op(ac, T._v(r))
        assert T.A[i] == ac
        return ac

Back to top page