"""Segment tree with lazy propagation: range add, range sum. Recursive implementation over a power-of-two padded array. The padding is load bearing: every internal node then spans an even length and splits into two equal halves, which is exactly what ``_push``'s child-length arithmetic (``mid - lo + 1`` for the left child, ``hi - mid`` for the right) relies on. Supports: build, point query (via range sum of length 1), range add, range sum. Bounds contract: indices are 0-indexed and inclusive. Any endpoint outside ``[0, len(data) - 1]`` raises ``IndexError`` -- nothing is clamped, so a query can never reach a padded cell that holds no data. An empty tree therefore accepts no index at all: there is no element to point at. A reversed range (``l > r``) is empty: it contributes 0 and is not bounds-checked. """ class SegmentTree: def __init__(self, data): self.n = len(data) # pad to a power of two: every internal node then spans an even length # and splits equally, so _push's two child-length expressions agree. 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 _span_ok(self, l, r): """Bounds contract for a [l, r] operation. False means empty range.""" if l > r: return False if l < 0 or r >= self.n: raise IndexError( "range [%d, %d] out of bounds for %d element(s) " "(valid indices 0..%d)" % (l, r, self.n, self.n - 1) ) return True def _push(self, node, lo, hi): if self.lazy[node] != 0: mid = (lo + hi) // 2 left = 2 * node right = 2 * node + 1 # sizes: left covers [lo, mid] = mid - lo + 1, right covers # [mid + 1, hi] = hi - mid; equal because spans are even. 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). Raises IndexError out of bounds.""" if not self._span_ok(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). Raises IndexError out of bounds.""" if not self._span_ok(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)]