beam: Fenwick (BIT) with point update, prefix/range sums, lower_bound + tests
This commit is contained in:
beam committed
1 parent
bc5f3e9e1c
commit
26cdf8bf7e
2 files changed
+229
No files matched your search
@@ -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)
|
||||
@@ -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")
|
||||
Reference in new issue
Block a user