Source code for titan_pylib.data_structures.set.dynamic_fenwick_tree_multiset

 1from titan_pylib.data_structures.set.dynamic_fenwick_tree_set import (
 2    DynamicFenwickTreeSet,
 3)
 4from typing import Iterable
 5
 6

[docs]
 7class DynamicFenwickTreeMultiset(DynamicFenwickTreeSet):
 8
 9    def __init__(self, n: int, a: Iterable[int] = []) -> None:
10        super().__init__(n, a)
11

[docs]
12    def add(self, key: int, val: int = 1) -> None:
13        self._len += val
14        if key in self._cnt:
15            self._cnt[key] += val
16        else:
17            self._cnt[key] = val
18        self._fw.add(key, val)

19

[docs]
20    def discard(self, key: int, val: int = 1) -> bool:
21        if key not in self._cnt:
22            return False
23        cnt = self._cnt[key]
24        if val >= cnt:
25            self._len -= cnt
26            del self._cnt[key]
27            self._fw.add(key, -cnt)
28        else:
29            self._len -= val
30            self._cnt[key] -= val
31            self._fw.add(key, -val)
32        return True

33

[docs]
34    def discard_all(self, key: int) -> bool:
35        return self.discard(key, val=self.count(key))

36

[docs]
37    def count(self, key: int) -> int:
38        return self._cnt[key]

39

[docs]
40    def items(self) -> Iterable[tuple[int, int]]:
41        _iter = 0
42        while _iter < self._len:
43            res = self.__getitem__(_iter)
44            cnt = self.count(res)
45            _iter += cnt
46            yield res, cnt

47

[docs]
48    def show(self) -> None:
49        print("{" + ", ".join(f"{i[0]}: {i[1]}" for i in self.items()) + "}")

50
51    def __repr__(self):
52        return f"DynamicFenwickTreeMultiset({self})"