beam: segment tree with lazy propagation (range add/sum) + randomized property tests

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)]

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")