cp_library/ds/tree/segtree_cls.py

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

Required by

• cp_library/alg/tree/fast/hld_commutative_cls.py
• cp_library/alg/tree/fast/hld_monoid_cls.py

Verified with

• test/aoj/dsl/dsl_2_a_segtree.test.py
• test/library-checker/data-structure/point_set_range_composite.test.py
• test/library-checker/tree/vertex_add_path_sum_hld_commutative.test.py
• test/library-checker/tree/vertex_add_path_sum_hld_monoid.test.py
• test/library-checker/tree/vertex_set_path_composite.test.py

Code

import cp_library.__header__
from typing import Callable, Generic, Union
from cp_library.misc.typing import _T
import cp_library.ds.__header__
import cp_library.ds.tree.__header__

class SegTree(Generic[_T]):
    def __init__(seg, op: Callable[[_T, _T], _T], e: _T, v: Union[int, list[_T]]) -> None:
        if isinstance(v, int): v = [e] * v
        seg.op, seg.e, seg.n = op, e, (n := len(v))
        seg.log, seg.sz, seg.d = (log := (n-1).bit_length()+1), (sz := 1 << log), [e] * (sz << 1)
        for i in range(n): seg.d[sz + i] = v[i]
        for i in range(sz-1,0,-1): seg.d[i] = op(seg.d[i<<1], seg.d[i<<1|1])

    def set(seg, p: int, x: _T) -> None:
        seg.d[p := p + seg.sz], op = x, seg.op
        for _ in range(seg.log): seg.d[p:=p>>1] = op(seg.d[p:=p^(p&1)], seg.d[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

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from typing import Callable, Generic, Union
from typing import TypeVar
_T = TypeVar('T')



class SegTree(Generic[_T]):
    def __init__(seg, op: Callable[[_T, _T], _T], e: _T, v: Union[int, list[_T]]) -> None:
        if isinstance(v, int): v = [e] * v
        seg.op, seg.e, seg.n = op, e, (n := len(v))
        seg.log, seg.sz, seg.d = (log := (n-1).bit_length()+1), (sz := 1 << log), [e] * (sz << 1)
        for i in range(n): seg.d[sz + i] = v[i]
        for i in range(sz-1,0,-1): seg.d[i] = op(seg.d[i<<1], seg.d[i<<1|1])

    def set(seg, p: int, x: _T) -> None:
        seg.d[p := p + seg.sz], op = x, seg.op
        for _ in range(seg.log): seg.d[p:=p>>1] = op(seg.d[p:=p^(p&1)], seg.d[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

Back to top page