1from typing import Union, Iterable, Optional
2
3
[docs]
4class FenwickTree:
5 """FenwickTreeです。"""
6
7 def __init__(self, n_or_a: Union[Iterable[int], int]):
8 """構築します。
9 :math:`O(n)` です。
10
11 Args:
12 n_or_a (Union[Iterable[int], int]): `n_or_a` が `int` のとき、初期値 `0` 、長さ `n` で構築します。
13 `n_or_a` が `Iterable` のとき、初期値 `a` で構築します。
14 """
15 if isinstance(n_or_a, int):
16 self._size = n_or_a
17 self._tree = [0] * (self._size + 1)
18 else:
19 a = n_or_a if isinstance(n_or_a, list) else list(n_or_a)
20 _size = len(a)
21 _tree = [0] + a
22 for i in range(1, _size):
23 if i + (i & -i) <= _size:
24 _tree[i + (i & -i)] += _tree[i]
25 self._size = _size
26 self._tree = _tree
27 self._s = 1 << (self._size - 1).bit_length()
28
[docs]
29 def pref(self, r: int) -> int:
30 """区間 ``[0, r)`` の総和を返します。
31 :math:`O(\\log{n})` です。
32 """
33 assert (
34 0 <= r <= self._size
35 ), f"IndexError: {self.__class__.__name__}.pref({r}), n={self._size}"
36 ret, _tree = 0, self._tree
37 while r > 0:
38 ret += _tree[r]
39 r &= r - 1
40 return ret
41
[docs]
42 def suff(self, l: int) -> int:
43 """区間 ``[l, n)`` の総和を返します。
44 :math:`O(\\log{n})` です。
45 """
46 assert (
47 0 <= l < self._size
48 ), f"IndexError: {self.__class__.__name__}.suff({l}), n={self._size}"
49 return self.pref(self._size) - self.pref(l)
50
[docs]
51 def sum(self, l: int, r: int) -> int:
52 """区間 ``[l, r)`` の総和を返します。
53 :math:`O(\\log{n})` です。
54 """
55 assert (
56 0 <= l <= r <= self._size
57 ), f"IndexError: {self.__class__.__name__}.sum({l}, {r}), n={self._size}"
58 _tree = self._tree
59 res = 0
60 while r > l:
61 res += _tree[r]
62 r &= r - 1
63 while l > r:
64 res -= _tree[l]
65 l &= l - 1
66 return res
67
68 prod = sum
69
[docs]
70 def __getitem__(self, k: int) -> int:
71 """位置 ``k`` の要素を返します。
72 :math:`O(\\log{n})` です。
73 """
74 assert (
75 -self._size <= k < self._size
76 ), f"IndexError: {self.__class__.__name__}[{k}], n={self._size}"
77 if k < 0:
78 k += self._size
79 return self.sum(k, k + 1)
80
[docs]
81 def add(self, k: int, x: int) -> None:
82 """``k`` 番目の値に ``x`` を加えます。
83 :math:`O(\\log{n})` です。
84 """
85 assert (
86 0 <= k < self._size
87 ), f"IndexError: {self.__class__.__name__}.add({k}, {x}), n={self._size}"
88 k += 1
89 _tree = self._tree
90 while k <= self._size:
91 _tree[k] += x
92 k += k & -k
93
[docs]
94 def __setitem__(self, k: int, x: int):
95 """``k`` 番目の値を ``x`` に更新します。
96 :math:`O(\\log{n})` です。
97 """
98 assert (
99 -self._size <= k < self._size
100 ), f"IndexError: {self.__class__.__name__}[{k}] = {x}, n={self._size}"
101 if k < 0:
102 k += self._size
103 pre = self[k]
104 self.add(k, x - pre)
105
[docs]
106 def bisect_left(self, w: int) -> Optional[int]:
107 i, s, _size, _tree = 0, self._s, self._size, self._tree
108 while s:
109 if i + s <= _size and _tree[i + s] < w:
110 w -= _tree[i + s]
111 i += s
112 s >>= 1
113 return i if w else None
114
[docs]
115 def bisect_right(self, w: int) -> int:
116 i, s, _size, _tree = 0, self._s, self._size, self._tree
117 while s:
118 if i + s <= _size and _tree[i + s] <= w:
119 w -= _tree[i + s]
120 i += s
121 s >>= 1
122 return i
123
124 def _pop(self, k: int) -> int:
125 assert k >= 0
126 i, acc, s, _size, _tree = 0, 0, self._s, self._size, self._tree
127 while s:
128 if i + s <= _size:
129 if acc + _tree[i + s] <= k:
130 acc += _tree[i + s]
131 i += s
132 else:
133 _tree[i + s] -= 1
134 s >>= 1
135 return i
136
[docs]
137 def tolist(self) -> list[int]:
138 """リストにして返します。
139 :math:`O(n)` です。
140 """
141 sub = [self.pref(i) for i in range(self._size + 1)]
142 return [sub[i + 1] - sub[i] for i in range(self._size)]
143
[docs]
144 @staticmethod
145 def get_inversion_num(a: list[int], compress: bool = False) -> int:
146 inv = 0
147 if compress:
148 a_ = sorted(set(a))
149 z = {e: i for i, e in enumerate(a_)}
150 fw = FenwickTree(len(a_) + 1)
151 for i, e in enumerate(a):
152 inv += i - fw.pref(z[e] + 1)
153 fw.add(z[e], 1)
154 else:
155 fw = FenwickTree(len(a) + 1)
156 for i, e in enumerate(a):
157 inv += i - fw.pref(e + 1)
158 fw.add(e, 1)
159 return inv
160
161 def __str__(self):
162 return str(self.tolist())
163
164 def __repr__(self):
165 return f"{self.__class__.__name__}({self})"