順序付き集合 (Treap)
(DataStructure/SortedSet/SortedSetTreap.py)

View this file on GitHub
Last update: 2021-05-14 22:52:43+09:00

概要

Treap による順序付き集合。

使い方

SortedSetTreap()
空の順序付き集合を作成する。計算量 $O(1)$

__contains__(key: int) -> bool
要素 key が集合に属しているかどうかを返す。計算量 $\mathrm{expected}\ O(\log N)$
__len__() -> int
集合の大きさを返す。計算量 $O(1)$
__iter__() -> Iterator[int]
要素の小さい順のイテレータオブジェクトを返す。イテレーション $1$ 回の計算量 $\mathrm{expected}\ O(\log N)$。全 $N$ 要素のイテレーションの計算量 $O(N)$
add(key: int) -> bool
要素 key を集合に追加する。追加に成功した場合は True を、失敗した場合 (既に key が集合に属していた場合) は False を返す。計算量 $\mathrm{expected}\ O(\log N)$
remove(key: int) -> bool
集合から key を削除する。削除に成功した場合は True を、失敗した場合 (key が集合に属していなかった場合) は False を返す。計算量 $\mathrm{expected}\ O(\log N)$
iterate(lower: int) -> Iterator[int]
lower 以上の要素に対して、小さい順のイテレータオブジェクトを返す。イテレーション $1$ 回の計算量 $\mathrm{expected}\ O(\log N)$。全 $k$ 要素のイテレーションの計算量 $O(k + \log N)$
next_val(lower: int) -> int
lower 以上の最小の要素を返す。そのような要素が存在しない場合は None を返す。計算量 $\mathrm{expected}\ O(\log N)$
prev_val(upper: int) -> int
upper よりも小さい最大の要素を返す。そのような要素が存在しない場合は None を返す。計算量 $\mathrm{expected}\ O(\log N)$

Depends on

乱数生成 (Xorshift) (misc/xorshift.py)

Verified with

TestCase/AOJ/ITP2_7_C.Treap.test.py

Code

from misc.xorshift import xorshift32
from array import array


class SortedSetTreap:
    def __init__(self):
        self.root = -1
        self.pow_sz = 1
        self.rand32 = xorshift32()
        self.stack = array("i", [0])
        self.l = array("i", [-1])
        self.r = array("i", [-1])
        self.p = array("i", [-1])
        self.pris = array("I", [self.rand32()])
        self.keys = [0]

    def __contains__(self, key):
        return self._search(key)

    def __len__(self):
        return self.pow_sz - len(self.stack)

    def __iter__(self):
        def dfs(nd):
            if self.l[nd] != -1:
                yield from dfs(self.l[nd])
            yield self.keys[nd]
            if self.r[nd] != -1:
                yield from dfs(self.r[nd])
        if self.root != -1:
            yield from dfs(self.root)

    def _new_node(self, key):
        if self.stack:
            nd = self.stack.pop()
            self.keys[nd] = key
            return nd
        else:
            self.stack = array("i", [nd for nd in range(self.pow_sz, self.pow_sz << 1)])
            self.l += array("i", [-1] * self.pow_sz)
            self.r += array("i", [-1] * self.pow_sz)
            self.p += array("i", [-1] * self.pow_sz)
            self.pris += array("I", [self.rand32() for _ in range(self.pow_sz)])
            self.keys += [0] * self.pow_sz
            self.pow_sz <<= 1
            return self._new_node(key)

    def _search(self, key):
        nd = self.root
        while nd != -1:
            if key < self.keys[nd]: nd = self.l[nd]
            elif key > self.keys[nd]: nd = self.r[nd]
            else: return True
        return False

    def add(self, key):
        if self.root == -1:
            self.root = self._new_node(key)
            return True

        nd = self.root
        while True:
            if key < self.keys[nd]:
                if self.l[nd] == -1:
                    self.l[nd] = self._new_node(key)
                    self.p[self.l[nd]] = nd
                    nd = self.l[nd]
                    break
                nd = self.l[nd]
            elif key > self.keys[nd]:
                if self.r[nd] == -1:
                    self.r[nd] = self._new_node(key)
                    self.p[self.r[nd]] = nd
                    nd = self.r[nd]
                    break
                nd = self.r[nd]
            else:
                return False

        while self.p[nd] != -1 and self.pris[self.p[nd]] > self.pris[nd]:
            if self.r[self.p[nd]] == nd: self._rotl(self.p[nd])
            else: self._rotr(self.p[nd])
        return True

    def _rotl(self, nd):
        pv = self.r[nd]
        self.p[pv] = self.p[nd]
        if self.p[pv] != -1:
            if self.l[self.p[pv]] == nd: self.l[self.p[pv]] = pv
            else: self.r[self.p[pv]] = pv
        self.r[nd] = self.l[pv]
        if self.r[nd] != -1: self.p[self.r[nd]] = nd
        self.p[nd] = pv
        self.l[pv] = nd
        if nd == self.root: self.root = pv

    def _rotr(self, nd):
        pv = self.l[nd]
        self.p[pv] = self.p[nd]
        if self.p[pv] != -1:
            if self.r[self.p[pv]] == nd: self.r[self.p[pv]] = pv
            else: self.l[self.p[pv]] = pv
        self.l[nd] = self.r[pv]
        if self.l[nd] != -1: self.p[self.l[nd]] = nd
        self.p[nd] = pv
        self.r[pv] = nd
        if nd == self.root: self.root = pv

    def remove(self, key):
        if self.root == -1: return False

        nd = self.root
        while True:
            if nd == -1: return False
            elif key < self.keys[nd]: nd = self.l[nd]
            elif key > self.keys[nd]: nd = self.r[nd]
            else: break

        while self.l[nd] != -1 or self.r[nd] != -1:
            if self.l[nd] == -1: self._rotl(nd)
            elif self.r[nd] == -1: self._rotr(nd)
            elif self.pris[self.l[nd]] < self.pris[self.r[nd]]: self._rotr(nd)
            else: self._rotl(nd)

        if nd == self.root: self.root = -1
        elif self.l[self.p[nd]] == nd: self.l[self.p[nd]] = -1
        else: self.r[self.p[nd]] = -1
        self.stack.append(nd)
        return True

    def iterate(self, lower):
        def dfs(nd):
            if self.l[nd] != -1 and (self.keys[nd] > lower):
                yield from dfs(self.l[nd])
            if self.keys[nd] >= lower:
                yield self.keys[nd]
            if self.r[nd] != -1:
                yield from dfs(self.r[nd])
        if self.root != -1:
            yield from dfs(self.root)

    def next_val(self, lower):
        ret = None
        nd = self.root
        while nd != -1:
            if self.keys[nd] >= lower:
                ret = self.keys[nd]
                nd = self.l[nd]
            else: nd = self.r[nd]
        return ret

    def prev_val(self, upper):
        ret = None
        nd = self.root
        while nd != -1:
            if self.keys[nd] < upper:
                ret = self.keys[nd]
                nd = self.r[nd]
            else: nd = self.l[nd]
        return ret

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

順序付き集合 (Treap) (DataStructure/SortedSet/SortedSetTreap.py)

概要

使い方

Depends on

Verified with

Code

順序付き集合 (Treap)
(DataStructure/SortedSet/SortedSetTreap.py)