順序付き集合 (Binary Indexed Tree)
(DataStructure/SortedSet/SortedSetBIT.py)

• View this file on GitHub
• Last update: 2021-02-06 18:04:37+09:00

概要

Binary Indexed Tree による順序付き集合。集合に属する可能性のある要素を、あらかじめ与える必要があることに注意。

使い方

SortedSetBIT(cands: Iterable[int])
cands 内の要素を元として含むことが可能な、空の順序付き集合を作成する。n = len(cands) とすると、計算量 $O(n \log n)$

• __contains__(val: int) -> bool
  要素 val が集合に属しているかどうかを返す。計算量 $O(\log n)$

• __len__() -> int
  集合の大きさを返す。計算量 $O(1)$

• add(val: int) -> bool
  要素 val を集合に追加する。追加に成功した場合は True を、失敗した場合 (既に val が集合に属していた場合) は False を返す。val が cands の要素ではない場合は例外 KeyError が発生するので注意。計算量 $O(\log n)$

• remove(val: int) -> bool
  集合から val を削除する。削除に成功した場合は True を、失敗した場合 (val が集合に属していなかった場合) は False を返す。計算量 $O(\log n)$

• count(vl: int, vr: int) -> int
  集合内の vl 以上かつ vr 未満である要素の数を返す。計算量 $O(\log n)$

• kth_smallest(k: int) -> int
  集合内で k 番目 (0-indexed) に小さい要素を返す。計算量 $O(\log n)$

• kth_largest(k: int) -> int
  集合内で k 番目 (0-indexed) に大きい要素を返す。計算量 $O(\log n)$

• prev_val(upper: int) -> int
  集合内で upper よりも小さい最大の要素を返す。そのような要素が存在しない場合は None を返す。計算量 $O(\log n)$

• next_val(lower: int) -> int
  集合内で lower 以上の最小の要素を返す。そのような要素が存在しない場合は None を返す。計算量 $O(\log n)$

• dump() -> List[int]
  集合内の全ての要素を小さい順に返す。計算量 $O(n)$

Required by

• 順序付き多重集合 (Binary Indexed Tree) (DataStructure/SortedSet/SortedMultiSetBIT.py)

Verified with

• TestCase/AOJ/ITP2_7_B.BIT.test.py

Code

from bisect import bisect_left


class SortedSetBIT:
    def __init__(self, cands):
        self.array = sorted(set(cands))
        self.comp = {val: i for i, val in enumerate(self.array)}
        self.size = len(self.array)
        self.cnt = 0
        self.bit = [0] * (self.size + 1)

    def __contains__(self, val):
        return self.count(val, val + 1) > 0

    def __len__(self):
        return self.cnt

    def _count(self, i):
        res = 0
        while i > 0:
            res += self.bit[i]
            i -= i & -i
        return res

    def add(self, val):
        if val in self:
            return False
        i = self.comp[val] + 1
        while i <= self.size:
            self.bit[i] += 1
            i += i & -i
        self.cnt += 1
        return True

    def remove(self, val):
        if val not in self:
            return False
        i = self.comp[val] + 1
        while i <= self.size:
            self.bit[i] -= 1
            i += i & -i
        self.cnt -= 1
        return True

    def count(self, vl, vr):
        l = bisect_left(self.array, vl)
        r = bisect_left(self.array, vr)
        return self._count(r) - self._count(l)

    def kth_smallest(self, k):
        if not(0 <= k < self.cnt):
            raise IndexError
        idx = 0
        k += 1
        d = 1 << self.size.bit_length()
        while d:
            if idx + d <= self.size and self.bit[idx + d] < k:
                k -= self.bit[idx + d]
                idx += d
            d >>= 1
        return self.array[idx]

    def kth_largest(self, k):
        return self.kth_smallest(self.cnt - k - 1)

    def prev_val(self, upper):
        upper = bisect_left(self.array, upper)
        k = self._count(upper) - 1
        return None if k == -1 else self.kth_smallest(k)

    def next_val(self, lower):
        lower = bisect_left(self.array, lower)
        k = self._count(lower)
        return None if k == self.cnt else self.kth_smallest(k)

    def all_dump(self):
        res = self.bit[:]
        for i in reversed(range(1, self.size)):
            if i + (i & -i) > self.size:
                continue
            res[i + (i & -i)] -= res[i]
        return [self.array[i] for i, flag in enumerate(res[1:]) if flag]

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/python.py", line 96, in bundle
    raise NotImplementedError
NotImplementedError

Back to top page