47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
"""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) |