Source code for titan_pylib.data_structures.set.set_quadratic_division

 1from typing import Iterable
 2
 3

[docs]
 4class SetQuadraticDivision:
 5
 6    # OrderedSet[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) -> bool:
21        if self.data[key]:
22            return False
23        self._len += 1
24        self.data[key] += 1
25        self.bucket_data[key // self.size] += 1
26        return True

27

[docs]
28    def discard(self, key: int, cnt) -> bool:
29        if not self.data[key]:
30            return False
31        self._len -= 1
32        self.data[key] -= cnt
33        self.bucket_data[key // self.size] -= cnt
34        return True

35

[docs]
36    def remove(self, key: int, cnt: int = 1) -> None:
37        self.data[key] -= cnt
38        self.bucket_data[key // self.size] -= cnt

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