97 lines
3.7 KiB
Python
97 lines
3.7 KiB
Python
"""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)]
|