Source code for titan_pylib.data_structures.fenwick_tree.dynamic_fenwick_tree

 1from typing import Optional, Final
 2
 3

[docs]
 4class DynamicFenwickTree:
 5    """必要なところだけノードを作ります。"""
 6
 7    def __init__(self, u: int):
 8        """Build DynamicFenwickTree [0, u)."""
 9        assert isinstance(
10            u, int
11        ), f"TypeError: DynamicFenwickTree({u}), {u} must be int"
12        self._u: Final[int] = u
13        self._tree: dict[int, int] = {}
14        self._s: Final[int] = 1 << (u - 1).bit_length()
15

[docs]
16    def add(self, k: int, x: int) -> None:
17        assert (
18            0 <= k < self._u
19        ), f"IndexError: DynamicFenwickTree.add({k}, {x}), u={self._u}"
20        k += 1
21        while k <= self._u:
22            if k in self._tree:
23                self._tree[k] += x
24            else:
25                self._tree[k] = x
26            k += k & -k

27

[docs]
28    def pref(self, r: int) -> int:
29        assert (
30            0 <= r <= self._u
31        ), f"IndexError: DynamicFenwickTree.pref({r}), u={self._u}"
32        ret = 0
33        while r > 0:
34            ret += self._tree.get(r, 0)
35            r -= r & -r
36        return ret

37

[docs]
38    def sum(self, l: int, r: int) -> int:
39        assert (
40            0 <= l <= r <= self._u
41        ), f"IndexError: DynamicFenwickTree.sum({l}, {r}), u={self._u}"
42        # return self.pref(r) - self.pref(l)
43        _tree = self._tree
44        res = 0
45        while r > l:
46            res += _tree.get(r, 0)
47            r -= r & -r
48        while l > r:
49            res -= _tree.get(l, 0)
50            l -= l & -l
51        return res

52

[docs]
53    def bisect_left(self, w: int) -> Optional[int]:
54        i, s = 0, self._s
55        while s:
56            if i + s <= self._u:
57                if i + s in self._tree and self._tree[i + s] < w:
58                    w -= self._tree[i + s]
59                    i += s
60                elif i + s not in self._tree and 0 < w:
61                    i += s
62            s >>= 1
63        return i if w else None

64

[docs]
65    def bisect_right(self, w: int) -> int:
66        i, s = 0, self._s
67        while s:
68            if i + s <= self._u:
69                if i + s in self._tree and self._tree[i + s] <= w:
70                    w -= self._tree[i + s]
71                    i += s
72                elif i + s not in self._tree and 0 <= w:
73                    i += s
74            s >>= 1
75        return i

76
77    def __str__(self):
78        return str(self._tree)