multiset_quadratic_division¶
ソースコード¶
from titan_pylib.data_structures.set.multiset_quadratic_division import MultisetQuadraticDivision
展開済みコード¶
1# from titan_pylib.data_structures.set.multiset_quadratic_division import MultisetQuadraticDivision
2from typing import Iterable
3
4
5class MultisetQuadraticDivision:
6
7 # OrderedMultiset[int]
8 # Space Complexity : O(u)
9 # add / discard / remove / contains : O(1)
10 # kth_elm : O(√u)
11
12 def __init__(self, u: int, a: Iterable[int] = []):
13 self.data = [0] * u
14 self.size = int(u**0.5) + 1
15 self.bucket_cnt = (u + self.size - 1) // self.size
16 self.bucket_data = [0] * self.bucket_cnt
17 self._len = 0
18 for e in a:
19 self.add(e)
20
21 def add(self, key: int, cnt: int = 1) -> None:
22 self._len += cnt
23 self.data[key] += cnt
24 self.bucket_data[key // self.size] += cnt
25
26 def discard(self, key: int, cnt: int = 1) -> None:
27 cnt = min(cnt, self.data[key])
28 self._len -= cnt
29 self.data[key] -= cnt
30 self.bucket_data[key // self.size] -= cnt
31
32 def remove(self, key: int, cnt: int = 1) -> None:
33 self.data[key] -= cnt
34 self.bucket_data[key // self.size] -= cnt
35
36 def count(self, key: int) -> int:
37 return self.data[key]
38
39 def __contains__(self, key: int):
40 return self.data[key] != 0
41
42 def __getitem__(self, k: int) -> int:
43 indx = 0
44 data, bucket_data = self.data, self.bucket_data
45 while bucket_data[indx] < k:
46 k -= bucket_data[indx]
47 indx += 1
48 indx *= self.size
49 while data[indx] == 0 or data[indx] <= k:
50 k -= data[indx]
51 indx += 1
52 return indx
53
54 def __len__(self):
55 return self._len