Source code for titan_pylib.data_structures.set.mex_multiset

 1from titan_pylib.data_structures.segment_tree.segment_tree import SegmentTree
 2from typing import Iterable
 3
 4

[docs]
 5class MexMultiset:
 6    """``MexMultiset`` です。
 7
 8    各操作は `log` がつきますが、ANDセグ木の ``log`` で割と軽いです。
 9    """
10
11    def __init__(self, u: int, a: Iterable[int] = []) -> None:
12        """``[0, u)`` の範囲の mex を計算する ``MexMultiset`` を構築します。
13
14        時間・空間共に :math:`O(u)` です。
15
16        Args:
17          u (int): 値の上限です。
18        """
19        data = [0] * (u + 1)
20        init_data = [1] * (u + 1)
21        for e in a:
22            if e <= u:
23                data[e] += 1
24                init_data[e] = 0
25        self.u: int = u
26        self.data: list[int] = data
27        self.seg: SegmentTree[int] = SegmentTree(init_data, op=lambda s, t: s | t, e=0)
28

[docs]
29    def add(self, key: int) -> None:
30        """``key`` を追加します。
31
32        :math:`O(\\log{n})` です。
33        """
34        if key > self.u:
35            return
36        if self.data[key] == 0:
37            self.seg[key] = 0
38        self.data[key] += 1

39

[docs]
40    def remove(self, key: int) -> None:
41        """``key`` を削除します。 ``key`` は存在していなければなりません。
42
43        :math:`O(\\log{n})` です。
44        """
45        if key > self.u:
46            return
47        if self.data[key] == 1:
48            self.seg[key] = 1
49        self.data[key] -= 1

50

[docs]
51    def mex(self) -> int:
52        """mex を返します。
53
54        :math:`O(\\log{n})` です。
55        """
56        return self.seg.max_right(0, lambda lr: lr == 0)