Fenwick tree in Python
Compact array structure for prefix aggregates with point updates in O(log n).
- File
fenwick_tree.py, 72 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/fenwick-tree/. The authors are listed in prov/T604D3.jsonl.
def lowbit(i: int) -> int:
"""The value of the lowest set bit of i, for i >= 1."""
return i & -i
class FenwickTree:
"""Prefix sums over A[1..n] of integers with point updates.
The positions are 1-based, as on the page. Each operation takes
O(log n) time; build takes O(n) time. The extra space is n integers.
"""
def __init__(self, n: int = 0) -> None:
"""Make the tree of an array of n zeros."""
self.n = n
self.t = [0] * (n + 1)
@classmethod
def build(cls, a: list[int]) -> "FenwickTree":
"""Make the tree of A[1..n] = a[0..n-1] in O(n) time."""
f = cls(len(a))
f.t[1:] = list(a) # @line L10
for i in range(1, f.n + 1): # @line L11
j = i + lowbit(i) # @line L12
if j <= f.n: # @line L13
f.t[j] += f.t[i]
return f
def add(self, i: int, delta: int) -> None:
"""Set A[i] <- A[i] + delta, for 1 <= i <= n."""
if not 1 <= i <= self.n:
raise IndexError(i)
while i <= self.n: # @line L1
self.t[i] += delta # @line L2
i += lowbit(i) # @line L3
def prefix(self, i: int) -> int:
"""Return A[1] + ... + A[i], for 0 <= i <= n."""
if not 0 <= i <= self.n:
raise IndexError(i)
s = 0 # @line L4
while i > 0: # @line L5
s += self.t[i] # @line L6
i -= lowbit(i) # @line L7
return s # @line L8
def range_sum(self, l: int, r: int) -> int:
"""Return A[l] + ... + A[r], for 1 <= l <= r <= n."""
return self.prefix(r) - self.prefix(l - 1) # @line L9
def search(self, s: int) -> int:
"""Return the least i with prefix(i) >= s, or n + 1 if there is none.
Precondition: every A[k] is non-negative.
"""
pos = 0 # @line L14
step = 1 << (self.n.bit_length() - 1) if self.n > 0 else 0
while step > 0: # @line L15
if pos + step <= self.n and self.t[pos + step] < s: # @line L16
pos += step # @line L17
s -= self.t[pos]
step //= 2 # @line L18
return pos + 1 # @line L19
def check_invariant(self, a: list[int]) -> None:
"""Assert that T[i] is the sum of A over (i - lowbit(i), i]; a is 0-based."""
for i in range(1, self.n + 1):
assert self.t[i] == sum(a[i - lowbit(i) : i])
test_fenwick_tree.py 73 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/fenwick-tree/. The authors are listed in prov/T604D3.jsonl.
import random
import pytest
from fenwick_tree import FenwickTree
def test_example() -> None:
a = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3]
f = FenwickTree.build(a)
assert f.t[1:] == [3, 4, 4, 9, 5, 14, 2, 31, 5, 8]
assert f.prefix(7) == 25
assert f.prefix(10) == 39
assert f.range_sum(3, 7) == 21
assert f.search(20) == 6
f.add(3, 2)
assert (f.t[3], f.t[4], f.t[8]) == (6, 11, 33)
assert f.range_sum(3, 7) == 23
def test_edges() -> None:
f = FenwickTree(0)
assert f.prefix(0) == 0
assert f.search(1) == 1
f = FenwickTree(1)
f.add(1, -5)
assert f.prefix(1) == -5 and f.range_sum(1, 1) == -5
with pytest.raises(IndexError):
f.add(0, 1)
big = 2**62
f = FenwickTree.build([big, big, big])
assert f.prefix(3) == 3 * big
f = FenwickTree.build([0, 0, 0, 0])
assert f.search(0) == 1 and f.search(1) == 5
def test_build_equals_adds() -> None:
rng = random.Random(7)
for n in range(40):
a = [rng.randint(-9, 9) for _ in range(n)]
g = FenwickTree(n)
for i, x in enumerate(a, 1):
g.add(i, x)
assert FenwickTree.build(a).t == g.t
def test_random_against_list() -> None:
rng = random.Random(1994)
for _ in range(500):
n = rng.randint(1, 40)
a = [rng.randint(0, 20) for _ in range(n)]
f = FenwickTree.build(a)
for _ in range(30):
op = rng.randrange(4)
if op == 0:
i, d = rng.randint(1, n), rng.randint(0, 9)
f.add(i, d)
a[i - 1] += d
elif op == 1:
l = rng.randint(1, n)
r = rng.randint(l, n)
assert f.range_sum(l, r) == sum(a[l - 1 : r])
elif op == 2:
i = rng.randint(0, n)
assert f.prefix(i) == sum(a[:i])
else:
s = rng.randint(0, sum(a) + 3)
want = next((i for i in range(1, n + 1) if sum(a[:i]) >= s), n + 1)
assert f.search(s) == want
f.check_invariant(a)
What the next level needs
- the tests did not run in this build