diff --git a/wizard-work/beam/fenwick.py b/wizard-work/beam/fenwick.py new file mode 100644 index 0000000..8d7b1df --- /dev/null +++ b/wizard-work/beam/fenwick.py @@ -0,0 +1,91 @@ +"""Fenwick tree (Binary Indexed Tree) over a 1-indexed array. + +Supports point updates, prefix sums, range sums, and lower_bound +(smallest index i such that prefix sum >= target). All operations O(log n). + +Invariants: +- tree[i] stores the sum of the range (i - lowbit(i) + 1 .. i). +- lowbit(x) = x & -x (two's complement). +- Updates propagate upward: i += lowbit(i). +- Queries accumulate downward: i -= lowbit(i). +""" +from __future__ import annotations + + +class Fenwick: + def __init__(self, n: int): + if n < 0: + raise ValueError("n must be >= 0") + self.n = n + self.tree = [0] * (n + 1) + + @classmethod + def from_array(cls, arr): + """Build in O(n) from a 0-indexed list.""" + n = len(arr) + ft = cls(n) + # copy + for i, v in enumerate(arr, start=1): + ft.tree[i] = v + # build: for each i, push to parent + for i in range(1, n + 1): + j = i + (i & -i) + if j <= n: + ft.tree[j] += ft.tree[i] + return ft + + def update(self, i: int, delta: int) -> None: + """Add delta at 1-indexed position i.""" + if i < 1 or i > self.n: + raise IndexError(f"index {i} out of range [1, {self.n}]") + while i <= self.n: + self.tree[i] += delta + i += i & -i + + def set(self, i: int, value: int) -> None: + """Set 1-indexed position i to value (requires current value tracking by caller? uses delta).""" + # Caller must know old; provide prefix-based? We'll implement via delta using a stored array. + raise NotImplementedError("use update with delta; or track base array externally") + + def prefix(self, i: int) -> int: + """Sum of positions 1..i (i may be 0).""" + if i < 0 or i > self.n: + raise IndexError(f"index {i} out of range [0, {self.n}]") + s = 0 + while i > 0: + s += self.tree[i] + i -= i & -i + return s + + def range_sum(self, l: int, r: int) -> int: + """Sum of positions l..r inclusive (1-indexed).""" + if l < 1 or r > self.n or l > r: + raise ValueError(f"invalid range [{l}, {r}]") + return self.prefix(r) - self.prefix(l - 1) + + def lower_bound(self, target: int) -> int: + """Smallest index i (1..n) such that prefix(i) >= target. + + Assumes all values non-negative and target >= 1. Returns n+1 if none. + """ + if target <= 0: + return 1 if self.n > 0 else 0 + pos = 0 + rem = target + # highest power of two <= n + bit = 1 + while bit * 2 <= self.n: + bit *= 2 + while bit > 0: + nxt = pos + bit + if nxt <= self.n and self.tree[nxt] < rem: + pos = nxt + rem -= self.tree[nxt] + bit >>= 1 + # pos is largest index with prefix < target + if pos + 1 <= self.n: + return pos + 1 + return self.n + 1 + + def total(self) -> int: + return self.prefix(self.n) \ No newline at end of file diff --git a/wizard-work/beam/test_fenwick.py b/wizard-work/beam/test_fenwick.py new file mode 100644 index 0000000..b7b4dad --- /dev/null +++ b/wizard-work/beam/test_fenwick.py @@ -0,0 +1,138 @@ +import random +from fenwick import Fenwick + + +def test_basic_update_prefix(): + ft = Fenwick(5) + for i in range(1, 6): + ft.update(i, i) + assert ft.prefix(0) == 0 + assert ft.prefix(1) == 1 + assert ft.prefix(3) == 6 + assert ft.prefix(5) == 15 + assert ft.total() == 15 + + +def test_range_sum(): + ft = Fenwick(6) + arr = [3, 1, 4, 1, 5, 9] + for i, v in enumerate(arr, start=1): + ft.update(i, v) + assert ft.range_sum(2, 4) == 1 + 4 + 1 + assert ft.range_sum(1, 6) == sum(arr) + assert ft.range_sum(3, 3) == 4 + + +def test_from_array(): + arr = [2, 5, 1, 3, 7] + ft = Fenwick.from_array(arr) + for i in range(1, len(arr) + 1): + assert ft.prefix(i) == sum(arr[:i]) + # update works on built tree + ft.update(3, 10) + assert ft.prefix(3) == 2 + 5 + 1 + 10 + + +def test_lower_bound(): + ft = Fenwick(5) + for i in range(1, 6): + ft.update(i, 1) # all ones + # prefix: 1,2,3,4,5 + assert ft.lower_bound(1) == 1 + assert ft.lower_bound(2) == 2 + assert ft.lower_bound(5) == 5 + assert ft.lower_bound(6) == 6 # none + + +def test_lower_bound_weighted(): + ft = Fenwick(4) + ft.update(1, 5) + ft.update(2, 2) + ft.update(3, 7) + ft.update(4, 1) + # prefixes: 5,7,14,15 + assert ft.lower_bound(5) == 1 + assert ft.lower_bound(6) == 2 + assert ft.lower_bound(7) == 2 + assert ft.lower_bound(8) == 3 + assert ft.lower_bound(15) == 4 + assert ft.lower_bound(16) == 5 + + +def test_random_vs_brute(): + n = 50 + arr = [0] * n + ft = Fenwick(n) + rng = random.Random(42) + for _ in range(2000): + op = rng.random() + if op < 0.5: + i = rng.randrange(1, n + 1) + d = rng.randrange(-5, 6) + arr[i - 1] += d + ft.update(i, d) + else: + # check a random prefix + i = rng.randrange(0, n + 1) + assert ft.prefix(i) == sum(arr[:i]), (i, ft.prefix(i), sum(arr[:i])) + # lower_bound against non-negative only: rebuild with non-negative + arr2 = [rng.randrange(0, 6) for _ in range(n)] + ft2 = Fenwick.from_array(arr2) + # verify lower_bound for every target 1..total + total = sum(arr2) + for t in range(1, total + 1): + # brute + s = 0 + brute = None + for idx, v in enumerate(arr2, start=1): + s += v + if s >= t: + brute = idx + break + if brute is None: + brute = n + 1 + assert ft2.lower_bound(t) == brute, (t, ft2.lower_bound(t), brute) + + +def test_bounds(): + ft = Fenwick(3) + try: + ft.update(0, 1) + except IndexError: + pass + else: + raise AssertionError("expected IndexError") + try: + ft.update(4, 1) + except IndexError: + pass + else: + raise AssertionError("expected IndexError") + try: + ft.prefix(4) + except IndexError: + pass + else: + raise AssertionError("expected IndexError") + try: + ft.range_sum(2, 1) + except ValueError: + pass + else: + raise AssertionError("expected ValueError") + + +def test_empty(): + ft = Fenwick(0) + assert ft.total() == 0 + assert ft.prefix(0) == 0 + + +if __name__ == "__main__": + # run all test_ functions + import sys + g = globals() + for name in sorted(g): + if name.startswith("test_"): + g[name]() + print("ALL TESTS PASSED")