diff --git a/wizard-work/beam/fenwick_range.py b/wizard-work/beam/fenwick_range.py new file mode 100644 index 0000000..c4613b6 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/wizard-work/beam/test_fenwick_range.py b/wizard-work/beam/test_fenwick_range.py new file mode 100644 index 0000000..b799619 --- /dev/null +++ b/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") \ No newline at end of file