beam: add FenwickRange (range add/range sum) with tests

wizard-work/beam/fenwick_range.py

@@ -0,0 +1,47 @@
+"""Fenwick tree supporting range add and range sum (1-indexed).
+
+Uses two BITs: B1, B2 such that prefix sum P(x) = sum_{i=1}^x a_i =
+    query(B1, x)*x - query(B2, x)
+where range add [l,r] += v does:
+    add(B1, l, v); add(B1, r+1, -v)
+    add(B2, l, v*(l-1)); add(B2, r+1, -v*r)
+"""
+from __future__ import annotations
+
+class FenwickRange:
+    def __init__(self, n: int):
+        if n < 0:
+            raise ValueError("n must be >= 0")
+        self.n = n
+        self.b1 = [0] * (n + 2)
+        self.b2 = [0] * (n + 2)
+
+    def _add(self, tree, i, delta):
+        while i <= self.n:
+            tree[i] += delta
+            i += i & -i
+
+    def _sum(self, tree, i):
+        s = 0
+        while i > 0:
+            s += tree[i]
+            i -= i & -i
+        return s
+
+    def range_add(self, l, r, v):
+        if l < 1 or r > self.n or l > r:
+            raise IndexError(f"range [{l},{r}] out of bounds [1,{self.n}]")
+        self._add(self.b1, l, v)
+        self._add(self.b1, r + 1, -v)
+        self._add(self.b2, l, v * (l - 1))
+        self._add(self.b2, r + 1, -v * r)
+
+    def prefix_sum(self, i):
+        if i < 0 or i > self.n:
+            raise IndexError(f"index {i} out of range [0,{self.n}]")
+        return self._sum(self.b1, i) * i - self._sum(self.b2, i)
+
+    def range_sum(self, l, r):
+        if l < 1 or r > self.n or l > r:
+            raise IndexError(f"range [{l},{r}] out of bounds [1,{self.n}]")
+        return self.prefix_sum(r) - self.prefix_sum(l - 1)

wizard-work/beam/test_fenwick_range.py

@@ -0,0 +1,64 @@
+import random
+from fenwick_range import FenwickRange
+
+def test_basic():
+    fr = FenwickRange(5)
+    fr.range_add(2, 4, 3)
+    assert fr.range_sum(1, 5) == 9
+    assert fr.range_sum(2, 4) == 9
+    assert fr.range_sum(1, 1) == 0
+    assert fr.range_sum(5, 5) == 0
+
+def test_multiple():
+    fr = FenwickRange(6)
+    fr.range_add(1, 6, 1)
+    fr.range_add(3, 5, 2)
+    arr = [1,1,3,3,3,1]
+    for i in range(1,7):
+        assert fr.prefix_sum(i) == sum(arr[:i])
+    assert fr.range_sum(2,4) == 1+3+3
+
+def test_random_vs_naive():
+    n=20
+    fr=FenwickRange(n)
+    arr=[0]*n
+    rng=random.Random(42)
+    for _ in range(200):
+        if rng.random()<0.5:
+            l=rng.randint(1,n); r=rng.randint(l,n); v=rng.randint(-5,5)
+            fr.range_add(l,r,v)
+            for i in range(l-1,r): arr[i]+=v
+        else:
+            l=rng.randint(1,n); r=rng.randint(l,n)
+            assert fr.range_sum(l,r)==sum(arr[l-1:r]), (l,r)
+    # prefix
+    for i in range(1,n+1):
+        assert fr.prefix_sum(i)==sum(arr[:i])
+
+def test_bounds():
+    fr=FenwickRange(3)
+    try:
+        fr.range_add(0,2,1); raise SystemExit("no raise")
+    except IndexError: pass
+    try:
+        fr.range_add(1,4,1); raise SystemExit("no raise")
+    except IndexError: pass
+    try:
+        fr.range_add(2,1,1); raise SystemExit("no raise")
+    except IndexError: pass
+    try:
+        fr.range_sum(1,4); raise SystemExit("no raise")
+    except IndexError: pass
+    try:
+        fr.prefix_sum(4); raise SystemExit("no raise")
+    except IndexError: pass
+
+def test_empty():
+    fr=FenwickRange(0)
+    # no operations valid
+    pass
+
+if __name__=="__main__":
+    for name in [n for n in globals() if n.startswith("test_")]:
+        globals()[name]()
+    print("ALL PASS")