91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
"""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) |