diff --git a/wizard-work/beam/segment_tree.py b/wizard-work/beam/segment_tree.py new file mode 100644 index 0000000..847e418 --- /dev/null +++ b/wizard-work/beam/segment_tree.py @@ -0,0 +1,70 @@ +"""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)] diff --git a/wizard-work/beam/test_segment_tree.py b/wizard-work/beam/test_segment_tree.py new file mode 100644 index 0000000..9a13424 --- /dev/null +++ b/wizard-work/beam/test_segment_tree.py @@ -0,0 +1,30 @@ +import random +from segment_tree import SegmentTree + +random.seed(42) +for trial in range(2000): + n = random.randint(1, 12) + data = [random.randint(-5, 5) for _ in range(n)] + st = SegmentTree(data) + arr = list(data) + for op in range(50): + if random.random() < 0.5: + l = random.randint(0, n-1) + r = random.randint(l, n-1) + v = random.randint(-5, 5) + st.range_add(l, r, v) + for i in range(l, r+1): + arr[i] += v + else: + l = random.randint(0, n-1) + r = random.randint(l, n-1) + got = st.range_sum(l, r) + exp = sum(arr[l:r+1]) + if got != exp: + print("MISMATCH trial", trial, "op", op, l, r, got, exp) + raise SystemExit(1) + # full list check + if st.to_list() != arr: + print("LIST MISMATCH", st.to_list(), arr) + raise SystemExit(1) +print("ALL 2000 TRIALS PASSED")