Extends Segment Tree to support range updates in O(log n) by deferring updates to children until needed.
Build: O(n) | Range Update/Query: O(log n)
class LazySegTree:
def __init__(self, n):
self.n = n
self.tree = [0] * (4 * n)
self.lazy = [0] * (4 * n)
def push_down(self, node, l, r):
if self.lazy[node]:
mid = (l + r) // 2
for child, size in [(2*node, mid-l+1), (2*node+1, r-mid)]:
self.tree[child] += self.lazy[node] * size
self.lazy[child] += self.lazy[node]
self.lazy[node] = 0
def update(self, node, l, r, ql, qr, val):
if ql <= l and r <= qr:
self.tree[node] += val * (r - l + 1)
self.lazy[node] += val
return
self.push_down(node, l, r)
mid = (l + r) // 2
if ql <= mid: self.update(2*node, l, mid, ql, qr, val)
if qr > mid: self.update(2*node+1, mid+1, r, ql, qr, val)
self.tree[node] = self.tree[2*node] + self.tree[2*node+1]
def query(self, node, l, r, ql, qr):
if ql <= l and r <= qr:
return self.tree[node]
self.push_down(node, l, r)
mid = (l + r) // 2
res = 0
if ql <= mid: res += self.query(2*node, l, mid, ql, qr)
if qr > mid: res += self.query(2*node+1, mid+1, r, ql, qr)
return res