Source code for titan_pylib.data_structures.set.dynamic_fenwick_tree_set

  1from titan_pylib.data_structures.fenwick_tree.dynamic_fenwick_tree import (
  2    DynamicFenwickTree,
  3)
  4from typing import Iterable, Optional
  5
  6

[docs]
  7class DynamicFenwickTreeSet:
  8
  9    # 整数[0, u)を、空間O(qlogu)/時間O(qlogu)
 10
 11    def __init__(self, u: int, a: Iterable[int] = []):
 12        self._size: int = u
 13        self._len: int = 0
 14        self._cnt: dict[int, int] = {}
 15        self._fw = DynamicFenwickTree(self._size)
 16        for _a in a:
 17            self.add(_a)
 18

[docs]
 19    def add(self, key: int) -> bool:
 20        if key in self._cnt:
 21            return False
 22        self._len += 1
 23        self._cnt[key] = 1
 24        self._fw.add(key, 1)
 25        return True

 26

[docs]
 27    def remove(self, key: int) -> None:
 28        if self.discard(key):
 29            return
 30        raise KeyError(key)

 31

[docs]
 32    def discard(self, key: int) -> bool:
 33        if key in self._cnt:
 34            self._len -= 1
 35            del self._cnt[key]
 36            self._fw.add(key, -1)
 37            return True
 38        return False

 39

[docs]
 40    def le(self, key: int) -> Optional[int]:
 41        if key in self._cnt:
 42            return key
 43        pref = self._fw.pref(key) - 1
 44        return None if pref < 0 else self._fw.bisect_right(pref)

 45

[docs]
 46    def lt(self, key: int) -> Optional[int]:
 47        pref = self._fw.pref(key) - 1
 48        return None if pref < 0 else self._fw.bisect_right(pref)

 49

[docs]
 50    def ge(self, key: int) -> Optional[int]:
 51        if key in self._cnt:
 52            return key
 53        pref = self._fw.pref(key + 1)
 54        return None if pref >= self._len else self._fw.bisect_right(pref)

 55

[docs]
 56    def gt(self, key: int) -> Optional[int]:
 57        pref = self._fw.pref(key + 1)
 58        return None if pref >= self._len else self._fw.bisect_right(pref)

 59

[docs]
 60    def index(self, key: int) -> int:
 61        return self._fw.pref(key)

 62

[docs]
 63    def index_right(self, key: int) -> int:
 64        return self._fw.pref(key + 1)

 65

[docs]
 66    def pop(self, k: int = -1) -> int:
 67        x = self.__getitem__(k)
 68        self.discard(x)
 69        return x

 70

[docs]
 71    def pop_max(self) -> int:
 72        assert self, "IndexError: pop_max() from empty DynamicFenwickTreeSet."
 73        return self.pop()

 74

[docs]
 75    def pop_min(self) -> int:
 76        assert self, "IndexError: pop_min() from empty DynamicFenwickTreeSet."
 77        return self.pop(0)

 78
 79    def __getitem__(self, k: int) -> int:
 80        if k < 0:
 81            k += self._len
 82        return self._fw.bisect_right(k)
 83
 84    def __iter__(self):
 85        self._iter = 0
 86        return self
 87
 88    def __next__(self):
 89        if self._iter == self._len:
 90            raise StopIteration
 91        res = self.__getitem__(self._iter)
 92        self._iter += 1
 93        return res
 94
 95    def __reversed__(self):
 96        for i in range(self._len):
 97            yield self.__getitem__(-i - 1)
 98
 99    def __len__(self):
100        return self._len
101
102    def __contains__(self, key: int):
103        return key in self._cnt
104
105    def __bool__(self):
106        return self._len > 0
107
108    def __str__(self):
109        return "{" + ", ".join(map(str, self)) + "}"
110
111    def __repr__(self):
112        return f"DynamicFenwickTreeSet({self})"