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