ikatakos
2/7/2019 - 11:10 AM

LazyBinaryIndexedTree for RangeAdd, RangeMax

[0, k) Add, [0, k) Max Query

class LazyBinaryTree:
    """ [0, k) Add, [0, k) Max """

    # For speeding up index calculation, [data] [lazy] are implemented 1-rooted
    # but 'update' and 'get' query requires 0-indexed

    # Not verified

    def __init__(self, n):
        n2 = 1 << (n - 1).bit_length()
        self.n = n2
        self.data = [0] * (n2 + 1)
        self.lazy = [0] * (n2 + 1)

    def _push(self, propagate):
        for i in reversed(propagate):
            v = self.lazy[i]
            if v == 0:
                continue
            j = i - 1
            while j:
                print('push', propagate, i, j)
                self.data[j] += v
                self.lazy[j] += v
                b = j & -j
                if j & (b << 1) == 0:
                    break
                j -= b
            self.lazy[i] = 0

    # [0, k)
    def add(self, k, x):
        i = k + 1
        mx = self.sentinel
        while i:
            self.data[i] += x
            self.lazy[i] += x
            mx = max(mx, self.data[i])
            i -= i & -i

    # [0, k)
    def get(self, k):
        i = k + 1
        propagate = []
        i += i & -i
        while i <= self.n:
            propagate.append(i)
            i += i & -i
        print('propagate', propagate)

        self._push(propagate)

        i = k + 1
        ret = 0
        while i:
            ret = max(ret, self.data[i])
            i -= i & -i

        tmp = self.data[k + 1]
        for i in propagate:
            tmp = self.data[i] = max(self.data[i], tmp)

        return ret

    def debug_print(self):
        print('data', self.data)
        print('lazy', self.lazy)