SortedMultiset

要素をソート済みの複数のバケットに分けて持つことで、C++ の std::multiset に近い操作を行います。要素は互いに比較できれば何でもよいです。

• SortedMultiset(a): 列 $a$ で初期化。 $O(N \log N)$
• add(x): $x$ を 1 つ追加
• discard(x): $x$ を 1 つ削除。削除できたら True
• x in s, count(x): 存在判定・個数
• lt(x), le(x), gt(x), ge(x): $x$ 未満の最大値・ $x$ 以下の最大値・ $x$ より大きい最小値・ $x$ 以上の最小値。存在しなければ None
• s[i]: 小さい方から $i$ 番目（0-indexed、負の添字も可）
• pop(i): $i$ 番目を削除して返す（既定は最大値）
• index(x), index_right(x): $x$ 未満・ $x$ 以下の要素の個数
• len(s), for x in s: 要素数・昇順の走査
• 計算量: 各操作 $O(\sqrt{N})$（初期化以外）

入っていない値を discard してもエラーにはならず、False が返ります。重複を許さない集合として使いたい場合は、add の前に x in s を確認します。

コード

import math
from bisect import bisect_left, bisect_right


class SortedMultiset:
    BUCKET_RATIO = 16
    SPLIT_RATIO = 24

    def __init__(self, a=[]):
        a = sorted(a)
        self.size = len(a)
        n = math.ceil(math.sqrt(self.size / self.BUCKET_RATIO))
        self.a = [a[self.size * i // n: self.size * (i + 1) // n] for i in range(n)]

    def __iter__(self):
        for b in self.a:
            yield from b

    def __reversed__(self):
        for b in reversed(self.a):
            yield from reversed(b)

    def __len__(self):
        return self.size

    def __repr__(self):
        return "SortedMultiset" + str(list(self))

    def _position(self, x):
        for i, b in enumerate(self.a):
            if x <= b[-1]:
                return b, i, bisect_left(b, x)
        return b, i, bisect_left(b, x)

    def __contains__(self, x):
        if self.size == 0:
            return False
        b, _, i = self._position(x)
        return i != len(b) and b[i] == x

    def count(self, x):
        return self.index_right(x) - self.index(x)

    def add(self, x):
        if self.size == 0:
            self.a = [[x]]
            self.size = 1
            return
        b, bi, i = self._position(x)
        b.insert(i, x)
        self.size += 1
        if len(b) > len(self.a) * self.SPLIT_RATIO:
            mid = len(b) >> 1
            self.a[bi:bi + 1] = [b[:mid], b[mid:]]

    def _pop(self, b, bi, i):
        x = b.pop(i)
        self.size -= 1
        if not b:
            del self.a[bi]
        return x

    def discard(self, x):
        if self.size == 0:
            return False
        b, bi, i = self._position(x)
        if i == len(b) or b[i] != x:
            return False
        self._pop(b, bi, i)
        return True

    def lt(self, x):
        for b in reversed(self.a):
            if b[0] < x:
                return b[bisect_left(b, x) - 1]

    def le(self, x):
        for b in reversed(self.a):
            if b[0] <= x:
                return b[bisect_right(b, x) - 1]

    def gt(self, x):
        for b in self.a:
            if b[-1] > x:
                return b[bisect_right(b, x)]

    def ge(self, x):
        for b in self.a:
            if b[-1] >= x:
                return b[bisect_left(b, x)]

    def __getitem__(self, i):
        if i < 0:
            for b in reversed(self.a):
                i += len(b)
                if i >= 0:
                    return b[i]
        else:
            for b in self.a:
                if i < len(b):
                    return b[i]
                i -= len(b)
        raise IndexError

    def pop(self, i=-1):
        if i < 0:
            for bi in range(len(self.a) - 1, -1, -1):
                b = self.a[bi]
                i += len(b)
                if i >= 0:
                    return self._pop(b, bi, i)
        else:
            for bi, b in enumerate(self.a):
                if i < len(b):
                    return self._pop(b, bi, i)
                i -= len(b)
        raise IndexError

    def index(self, x):
        res = 0
        for b in self.a:
            if b[-1] >= x:
                return res + bisect_left(b, x)
            res += len(b)
        return res

    def index_right(self, x):
        res = 0
        for b in self.a:
            if b[-1] > x:
                return res + bisect_right(b, x)
            res += len(b)
        return res