1from typing import Union, Iterable, TypeVar, Generic, Callable
2
3T = TypeVar("T")
4
5
[docs]
6class FenwickTreeAbst(Generic[T]):
7 """和や逆元をこちらで定義できます。"""
8
9 def __init__(
10 self,
11 n_or_a: Union[Iterable[T], T],
12 op: Callable[[T, T], T],
13 inv: Callable[[T], T],
14 e: T,
15 ) -> None:
16 if isinstance(n_or_a, int):
17 self._size = n_or_a
18 self._tree = [e] * (self._size + 1)
19 else:
20 a = n_or_a if isinstance(n_or_a, list) else list(n_or_a)
21 self._size = len(a)
22 self._tree = [e] + a
23 for i in range(1, self._size):
24 if i + (i & -i) <= self._size:
25 self._tree[i + (i & -i)] = op(
26 self._tree[i + (i & -i)], self._tree[i]
27 )
28 self.op = op
29 self.inv = inv
30 self.e = e
31 self._s = 1 << (self._size - 1).bit_length()
32
[docs]
33 def pref(self, r: int) -> T:
34 assert (
35 0 <= r <= self._size
36 ), f"IndexError: {self.__class__.__name__}.pref({r}), n={self._size}"
37 ret = self.e
38 while r > 0:
39 ret = self.op(ret, self._tree[r])
40 r -= r & -r
41 return ret
42
[docs]
43 def suff(self, l: int) -> T:
44 assert (
45 0 <= l < self._size
46 ), f"IndexError: {self.__class__.__name__}.suff({l}), n={self._size}"
47 return self.op(self.pref(self._size), self.inv(self.pref(l)))
48
[docs]
49 def sum(self, l: int, r: int) -> T:
50 assert (
51 0 <= l <= r <= self._size
52 ), f"IndexError: {self.__class__.__name__}.sum({l}, {r}), n={self._size}"
53 _tree = self._tree
54 res = self.e
55 while r > l:
56 res = self.op(res, _tree[r])
57 r &= r - 1
58 while l > r:
59 res += self.inv(_tree[l])
60 l &= l - 1
61 return res
62
63 prod = sum
64
65 def __getitem__(self, k: int) -> T:
66 assert (
67 -self._size <= k < self._size
68 ), f"IndexError: {self.__class__.__name__}.__getitem__({k}), n={self._size}"
69 if k < 0:
70 k += self._size
71 return self.op(self.pref(k + 1), self.inv(self.pref(k)))
72
[docs]
73 def add(self, k: int, x: T) -> None:
74 assert (
75 0 <= k < self._size
76 ), f"IndexError: {self.__class__.__name__}.add({k}, {x}), n={self._size}"
77 k += 1
78 while k <= self._size:
79 self._tree[k] = self.op(self._tree[k], x)
80 k += k & -k
81
82 def __setitem__(self, k: int, x: T):
83 assert (
84 -self._size <= k < self._size
85 ), f"IndexError: {self.__class__.__name__}.__setitem__({k}, {x}), n={self._size}"
86 if k < 0:
87 k += self._size
88 pre = self.__getitem__(k)
89 self.add(k, self.op(x, self.inv(pre)))
90
[docs]
91 def tolist(self) -> list[T]:
92 sub = [self.pref(i) for i in range(self._size + 1)]
93 return [self.op(sub[i + 1], self.inv(sub[i])) for i in range(self._size)]
94
95 def __str__(self):
96 return str(self.tolist())
97
98 def __repr__(self):
99 return f"{self.__class__.__name__}({self})"