1# from titan_pylib.data_structures.set.dynamic_fenwick_tree_set import DynamicFenwickTreeSet
2# from titan_pylib.data_structures.fenwick_tree.dynamic_fenwick_tree import (
3# DynamicFenwickTree,
4# )
5from typing import Optional, Final
6
7
8class DynamicFenwickTree:
9 """必要なところだけノードを作ります。"""
10
11 def __init__(self, u: int):
12 """Build DynamicFenwickTree [0, u)."""
13 assert isinstance(
14 u, int
15 ), f"TypeError: DynamicFenwickTree({u}), {u} must be int"
16 self._u: Final[int] = u
17 self._tree: dict[int, int] = {}
18 self._s: Final[int] = 1 << (u - 1).bit_length()
19
20 def add(self, k: int, x: int) -> None:
21 assert (
22 0 <= k < self._u
23 ), f"IndexError: DynamicFenwickTree.add({k}, {x}), u={self._u}"
24 k += 1
25 while k <= self._u:
26 if k in self._tree:
27 self._tree[k] += x
28 else:
29 self._tree[k] = x
30 k += k & -k
31
32 def pref(self, r: int) -> int:
33 assert (
34 0 <= r <= self._u
35 ), f"IndexError: DynamicFenwickTree.pref({r}), u={self._u}"
36 ret = 0
37 while r > 0:
38 ret += self._tree.get(r, 0)
39 r -= r & -r
40 return ret
41
42 def sum(self, l: int, r: int) -> int:
43 assert (
44 0 <= l <= r <= self._u
45 ), f"IndexError: DynamicFenwickTree.sum({l}, {r}), u={self._u}"
46 # return self.pref(r) - self.pref(l)
47 _tree = self._tree
48 res = 0
49 while r > l:
50 res += _tree.get(r, 0)
51 r -= r & -r
52 while l > r:
53 res -= _tree.get(l, 0)
54 l -= l & -l
55 return res
56
57 def bisect_left(self, w: int) -> Optional[int]:
58 i, s = 0, self._s
59 while s:
60 if i + s <= self._u:
61 if i + s in self._tree and self._tree[i + s] < w:
62 w -= self._tree[i + s]
63 i += s
64 elif i + s not in self._tree and 0 < w:
65 i += s
66 s >>= 1
67 return i if w else None
68
69 def bisect_right(self, w: int) -> int:
70 i, s = 0, self._s
71 while s:
72 if i + s <= self._u:
73 if i + s in self._tree and self._tree[i + s] <= w:
74 w -= self._tree[i + s]
75 i += s
76 elif i + s not in self._tree and 0 <= w:
77 i += s
78 s >>= 1
79 return i
80
81 def __str__(self):
82 return str(self._tree)
83from typing import Iterable, Optional
84
85
86class DynamicFenwickTreeSet:
87
88 # 整数[0, u)を、空間O(qlogu)/時間O(qlogu)
89
90 def __init__(self, u: int, a: Iterable[int] = []):
91 self._size: int = u
92 self._len: int = 0
93 self._cnt: dict[int, int] = {}
94 self._fw = DynamicFenwickTree(self._size)
95 for _a in a:
96 self.add(_a)
97
98 def add(self, key: int) -> bool:
99 if key in self._cnt:
100 return False
101 self._len += 1
102 self._cnt[key] = 1
103 self._fw.add(key, 1)
104 return True
105
106 def remove(self, key: int) -> None:
107 if self.discard(key):
108 return
109 raise KeyError(key)
110
111 def discard(self, key: int) -> bool:
112 if key in self._cnt:
113 self._len -= 1
114 del self._cnt[key]
115 self._fw.add(key, -1)
116 return True
117 return False
118
119 def le(self, key: int) -> Optional[int]:
120 if key in self._cnt:
121 return key
122 pref = self._fw.pref(key) - 1
123 return None if pref < 0 else self._fw.bisect_right(pref)
124
125 def lt(self, key: int) -> Optional[int]:
126 pref = self._fw.pref(key) - 1
127 return None if pref < 0 else self._fw.bisect_right(pref)
128
129 def ge(self, key: int) -> Optional[int]:
130 if key in self._cnt:
131 return key
132 pref = self._fw.pref(key + 1)
133 return None if pref >= self._len else self._fw.bisect_right(pref)
134
135 def gt(self, key: int) -> Optional[int]:
136 pref = self._fw.pref(key + 1)
137 return None if pref >= self._len else self._fw.bisect_right(pref)
138
139 def index(self, key: int) -> int:
140 return self._fw.pref(key)
141
142 def index_right(self, key: int) -> int:
143 return self._fw.pref(key + 1)
144
145 def pop(self, k: int = -1) -> int:
146 x = self.__getitem__(k)
147 self.discard(x)
148 return x
149
150 def pop_max(self) -> int:
151 assert self, "IndexError: pop_max() from empty DynamicFenwickTreeSet."
152 return self.pop()
153
154 def pop_min(self) -> int:
155 assert self, "IndexError: pop_min() from empty DynamicFenwickTreeSet."
156 return self.pop(0)
157
158 def __getitem__(self, k: int) -> int:
159 if k < 0:
160 k += self._len
161 return self._fw.bisect_right(k)
162
163 def __iter__(self):
164 self._iter = 0
165 return self
166
167 def __next__(self):
168 if self._iter == self._len:
169 raise StopIteration
170 res = self.__getitem__(self._iter)
171 self._iter += 1
172 return res
173
174 def __reversed__(self):
175 for i in range(self._len):
176 yield self.__getitem__(-i - 1)
177
178 def __len__(self):
179 return self._len
180
181 def __contains__(self, key: int):
182 return key in self._cnt
183
184 def __bool__(self):
185 return self._len > 0
186
187 def __str__(self):
188 return "{" + ", ".join(map(str, self)) + "}"
189
190 def __repr__(self):
191 return f"DynamicFenwickTreeSet({self})"