1# from titan_pylib.data_structures.fenwick_tree.dynamic_fenwick_tree import DynamicFenwickTree
2from typing import Optional, Final
3
4
5class DynamicFenwickTree:
6 """必要なところだけノードを作ります。"""
7
8 def __init__(self, u: int):
9 """Build DynamicFenwickTree [0, u)."""
10 assert isinstance(
11 u, int
12 ), f"TypeError: DynamicFenwickTree({u}), {u} must be int"
13 self._u: Final[int] = u
14 self._tree: dict[int, int] = {}
15 self._s: Final[int] = 1 << (u - 1).bit_length()
16
17 def add(self, k: int, x: int) -> None:
18 assert (
19 0 <= k < self._u
20 ), f"IndexError: DynamicFenwickTree.add({k}, {x}), u={self._u}"
21 k += 1
22 while k <= self._u:
23 if k in self._tree:
24 self._tree[k] += x
25 else:
26 self._tree[k] = x
27 k += k & -k
28
29 def pref(self, r: int) -> int:
30 assert (
31 0 <= r <= self._u
32 ), f"IndexError: DynamicFenwickTree.pref({r}), u={self._u}"
33 ret = 0
34 while r > 0:
35 ret += self._tree.get(r, 0)
36 r -= r & -r
37 return ret
38
39 def sum(self, l: int, r: int) -> int:
40 assert (
41 0 <= l <= r <= self._u
42 ), f"IndexError: DynamicFenwickTree.sum({l}, {r}), u={self._u}"
43 # return self.pref(r) - self.pref(l)
44 _tree = self._tree
45 res = 0
46 while r > l:
47 res += _tree.get(r, 0)
48 r -= r & -r
49 while l > r:
50 res -= _tree.get(l, 0)
51 l -= l & -l
52 return res
53
54 def bisect_left(self, w: int) -> Optional[int]:
55 i, s = 0, self._s
56 while s:
57 if i + s <= self._u:
58 if i + s in self._tree and self._tree[i + s] < w:
59 w -= self._tree[i + s]
60 i += s
61 elif i + s not in self._tree and 0 < w:
62 i += s
63 s >>= 1
64 return i if w else None
65
66 def bisect_right(self, w: int) -> int:
67 i, s = 0, self._s
68 while s:
69 if i + s <= self._u:
70 if i + s in self._tree and self._tree[i + s] <= w:
71 w -= self._tree[i + s]
72 i += s
73 elif i + s not in self._tree and 0 <= w:
74 i += s
75 s >>= 1
76 return i
77
78 def __str__(self):
79 return str(self._tree)