71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""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)]
|