セグメント木

2022年11月5日

説明

セグメント木

class SegTree:

    def __init__(self, n, operator=max, identity=-INFTY):
        '''
        1-indexedのセグメント木用配列を用意

        operator, identityの例

        operator: add
        identity: 0

        operator: mul
        identity: 1

        operator: min
        identity: INFTY

        operator: max
        identity: -INFTY

        operator: gcd
        identity: 1

        operator: xor
        identity: 0

        '''
        self.operator = operator
        self.identity = identity

        m = 1
        while 1<<m < n:
            m += 1
        self.siz = 1<<m
        self.data = [self.identity for _ in range(2*self.siz)]
        return

    def set(self, pos, x):
        '''
        pos: 1 始まりとしたときの値を変更する leaf の位置
        '''
        pos += self.siz - 1
        self.data[pos] = x
        pos //= 2
        while pos > 0:
            self.data[pos] = self.operator(self.data[2*pos], self.data[2*pos+1])
            pos //= 2
        return

    def get(self, pos):
        '''
        pos: 1 始まりとしたときの leaf の値を返す
        '''
        pos += self.siz - 1
        return self.data[pos]

    def query(self, l, r):
        '''
        1 始まりとしたときの半開区間 [l, r) での operator の結果を返す
        '''
        a = 1
        b = self.siz+1
        u = 1
        return self.query_(l, r, a, b, u)


    def query_(self, l, r, a, b, u):
        if r <= a or b <= l:
            return self.identity
        if l <= a and b <= r:
            return self.data[u]
        m = (a+b)//2
        ansl = self.query_(l, r, a, m, u*2)
        ansr = self.query_(l, r, m, b, 2*u+1)
        return self.operator(ansl, ansr)

    def print(self):
        print(self.data)