1# from titan_pylib.data_structures.set.fenwick_tree_set_fast import FenwickTreeSetFast
2from typing import Union, Iterable, Optional
3
4
5class FenwickTreeSetFast:
6
7 class InternalFenwickTree:
8
9 def __init__(self, _n_or_a: Union[int, list[int]]):
10 if isinstance(_n_or_a, int):
11 self._size = _n_or_a
12 self._tree = [0] * (self._size + 1)
13 else:
14 self._size = len(_n_or_a)
15 _tree = [0] + _n_or_a
16 for i in range(1, self._size):
17 if i + (i & -i) <= self._size:
18 _tree[i + (i & -i)] += _tree[i]
19 self._tree = _tree
20 self._s = 1 << (self._size - 1).bit_length()
21
22 def pref(self, r: int) -> int:
23 ret, _tree = 0, self._tree
24 while r > 0:
25 ret += _tree[r]
26 r &= r - 1
27 return ret
28
29 def add(self, k: int, x: int) -> None:
30 _size, _tree = self._size, self._tree
31 k += 1
32 while k <= _size:
33 _tree[k] += x
34 k += k & -k
35
36 def bisect_left(self, w: int) -> Optional[int]:
37 i, s, _size, _tree = 0, self._s, self._size, self._tree
38 while s:
39 if i + s <= _size and _tree[i + s] < w:
40 w -= _tree[i + s]
41 i += s
42 s >>= 1
43 return i if w else None
44
45 def bisect_right(self, w: int) -> int:
46 i, s, _size, _tree = 0, self._s, self._size, self._tree
47 while s:
48 if i + s <= _size and _tree[i + s] <= w:
49 w -= _tree[i + s]
50 i += s
51 s >>= 1
52 return i
53
54 def __init__(self, _u: int, a: Iterable[int] = []):
55 _len = 0
56 _size = _u
57 _cnt = [0] * _u
58 a_ = []
59 if a:
60 a_ = [0] * _u
61 for v in a:
62 assert (
63 v < _u
64 ), f"IndexError: FenwickTreeSetFast.__init__({_u}, {a}), not ({v} < {_u})"
65 if a_[v] == 0:
66 _len += 1
67 _cnt[v] = 1
68 a_[v] = 1
69 self._len = _len
70 self._cnt = _cnt
71 self._size = _size
72 self._fw = self.InternalFenwickTree(a_ if a else _u)
73
74 def add(self, key: int) -> bool:
75 if self._cnt[key]:
76 return False
77 self._len += 1
78 self._cnt[key] = 1
79 self._fw.add(key, 1)
80 return True
81
82 def remove(self, key: int) -> None:
83 if not self.discard(key):
84 raise KeyError(key)
85
86 def discard(self, key: int) -> bool:
87 if self._cnt[key]:
88 self._len -= 1
89 self._cnt[key] = 0
90 self._fw.add(key, -1)
91 return True
92 return False
93
94 def le(self, key: int) -> Optional[int]:
95 if self._cnt[key]:
96 return key
97 pref = self._fw.pref(key) - 1
98 return None if pref < 0 else self._fw.bisect_right(pref)
99
100 def lt(self, key: int) -> Optional[int]:
101 pref = self._fw.pref(key) - 1
102 return None if pref < 0 else self._fw.bisect_right(pref)
103
104 def ge(self, key: int) -> Optional[int]:
105 if self._cnt[key]:
106 return key
107 pref = self._fw.pref(key + 1)
108 return None if pref >= self._len else self._fw.bisect_right(pref)
109
110 def gt(self, key: int) -> Optional[int]:
111 pref = self._fw.pref(key + 1)
112 return None if pref >= self._len else self._fw.bisect_right(pref)
113
114 def index(self, key: int) -> int:
115 return self._fw.pref(key)
116
117 def index_right(self, key: int) -> int:
118 return self._fw.pref(key + 1)
119
120 def pop(self, k: int = -1) -> int:
121 if k < 0:
122 k += self._len
123 self._len -= 1
124 i, s, _size, _tree = 0, self._fw._s, self._fw._size, self._fw._tree
125 while s:
126 if i + s <= _size:
127 if _tree[i + s] <= k:
128 k -= _tree[i + s]
129 i += s
130 else:
131 _tree[i + s] -= 1
132 s >>= 1
133 self._cnt[i] = 0
134 return i
135
136 def pop_max(self) -> int:
137 return self.pop()
138
139 def pop_min(self) -> int:
140 return self.pop(0)
141
142 def __getitem__(self, k: int) -> int:
143 if k < 0:
144 k += self._len
145 return self._fw.bisect_right(k)
146
147 def __contains__(self, key: int):
148 return self._cnt[key]
149
150 def __str__(self):
151 return "{" + ", ".join(map(str, self)) + "}"
152
153 def __iter__(self):
154 self._iter = 0
155 return self
156
157 def __next__(self):
158 if self._iter == self._len:
159 raise StopIteration
160 res = self.__getitem__(self._iter)
161 self._iter += 1
162 return res
163
164 def __reversed__(self):
165 for i in range(self._len):
166 yield self.__getitem__(-i - 1)
167
168 def __len__(self):
169 return self._len
170
171 def __bool__(self):
172 return self._len > 0
173
174 def __repr__(self):
175 return f"FenwickTreeSetFast({self})"