set_quadratic_division¶
ソースコード¶
from titan_pylib.data_structures.set.set_quadratic_division import SetQuadraticDivision
展開済みコード¶
1# from titan_pylib.data_structures.set.set_quadratic_division import SetQuadraticDivision
2from typing import Iterable
3
4
5class SetQuadraticDivision:
6
7 # OrderedSet[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) -> bool:
22 if self.data[key]:
23 return False
24 self._len += 1
25 self.data[key] += 1
26 self.bucket_data[key // self.size] += 1
27 return True
28
29 def discard(self, key: int, cnt) -> bool:
30 if not self.data[key]:
31 return False
32 self._len -= 1
33 self.data[key] -= cnt
34 self.bucket_data[key // self.size] -= cnt
35 return True
36
37 def remove(self, key: int, cnt: int = 1) -> None:
38 self.data[key] -= cnt
39 self.bucket_data[key // self.size] -= cnt
40
41 def __contains__(self, key: int):
42 return self.data[key] != 0
43
44 def __getitem__(self, k: int) -> int:
45 indx = 0
46 data, bucket_data = self.data, self.bucket_data
47 while bucket_data[indx] < k:
48 k -= bucket_data[indx]
49 indx += 1
50 indx *= self.size
51 while data[indx] == 0 or data[indx] <= k:
52 k -= data[indx]
53 indx += 1
54 return indx
55
56 def __len__(self):
57 return self._len