"""Segment tree with lazy propagation: range add, range sum. Iterative-friendly recursive implementation. 1-indexed internally. Supports: build, point query (via range sum length 1), range add, range sum. """ class SegmentTree: def __init__(self, data): self.n = len(data) self.size = 1 while self.size < max(1, self.n): self.size <<= 1 self.tree = [0] * (2 * self.size) self.lazy = [0] * (2 * self.size) # build leaves for i in range(self.n): self.tree[self.size + i] = data[i] for i in range(self.size - 1, 0, -1): self.tree[i] = self.tree[2 * i] + self.tree[2 * i + 1] def _push(self, node, lo, hi): if self.lazy[node] != 0: mid = (lo + hi) // 2 left = 2 * node right = 2 * node + 1 self.tree[left] += self.lazy[node] * (mid - lo + 1) self.lazy[left] += self.lazy[node] self.tree[right] += self.lazy[node] * (hi - mid) self.lazy[right] += self.lazy[node] self.lazy[node] = 0 def _update(self, node, lo, hi, l, r, val): if r < lo or hi < l: return if l <= lo and hi <= r: self.tree[node] += val * (hi - lo + 1) self.lazy[node] += val return self._push(node, lo, hi) mid = (lo + hi) // 2 self._update(2 * node, lo, mid, l, r, val) self._update(2 * node + 1, mid + 1, hi, l, r, val) self.tree[node] = self.tree[2 * node] + self.tree[2 * node + 1] def range_add(self, l, r, val): """Add val to [l, r] inclusive (0-indexed).""" if l > r: return self._update(1, 0, self.size - 1, l, r, val) def _query(self, node, lo, hi, l, r): if r < lo or hi < l: return 0 if l <= lo and hi <= r: return self.tree[node] self._push(node, lo, hi) mid = (lo + hi) // 2 return self._query(2 * node, lo, mid, l, r) + self._query(2 * node + 1, mid + 1, hi, l, r) def range_sum(self, l, r): """Sum of [l, r] inclusive (0-indexed).""" if l > r: return 0 return self._query(1, 0, self.size - 1, l, r) def point_get(self, i): return self.range_sum(i, i) def to_list(self): return [self.point_get(i) for i in range(self.n)]