dynamic_fenwick_tree_multiset

ソースコード

from titan_pylib.data_structures.set.dynamic_fenwick_tree_multiset import DynamicFenwickTreeMultiset

view on github

展開済みコード

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

仕様

class DynamicFenwickTreeMultiset(n: int, a: Iterable[int] = [])[source]

Bases: DynamicFenwickTreeSet

add(key: int, val: int = 1) → None[source]
count(key: int) → int[source]
discard(key: int, val: int = 1) → bool[source]
discard_all(key: int) → bool[source]
items() → Iterable[tuple[int, int]][source]
show() → None[source]