Source code for titan_pylib.data_structures.set.multiset_quadratic_division

 1from typing import Iterable
 2
 3

[docs]
 4class MultisetQuadraticDivision:
 5
 6    # OrderedMultiset[int]
 7    # Space Complexity : O(u)
 8    # add / discard / remove / contains : O(1)
 9    # kth_elm : O(√u)
10
11    def __init__(self, u: int, a: Iterable[int] = []):
12        self.data = [0] * u
13        self.size = int(u**0.5) + 1
14        self.bucket_cnt = (u + self.size - 1) // self.size
15        self.bucket_data = [0] * self.bucket_cnt
16        self._len = 0
17        for e in a:
18            self.add(e)
19

[docs]
20    def add(self, key: int, cnt: int = 1) -> None:
21        self._len += cnt
22        self.data[key] += cnt
23        self.bucket_data[key // self.size] += cnt

24

[docs]
25    def discard(self, key: int, cnt: int = 1) -> None:
26        cnt = min(cnt, self.data[key])
27        self._len -= cnt
28        self.data[key] -= cnt
29        self.bucket_data[key // self.size] -= cnt

30

[docs]
31    def remove(self, key: int, cnt: int = 1) -> None:
32        self.data[key] -= cnt
33        self.bucket_data[key // self.size] -= cnt

34

[docs]
35    def count(self, key: int) -> int:
36        return self.data[key]

37
38    def __contains__(self, key: int):
39        return self.data[key] != 0
40
41    def __getitem__(self, k: int) -> int:
42        indx = 0
43        data, bucket_data = self.data, self.bucket_data
44        while bucket_data[indx] < k:
45            k -= bucket_data[indx]
46            indx += 1
47        indx *= self.size
48        while data[indx] == 0 or data[indx] <= k:
49            k -= data[indx]
50            indx += 1
51        return indx
52
53    def __len__(self):
54        return self._len