SortedMultiSet として使う BIT

Python は Hash Set 以外の集合データ構造を標準で実装しておらず貧弱なので順序が必要な場合は自前のデータ構造が必要になります。

万能ではありませんが、BIT でいくらか代用可能です。

BIT (Binary Indexed Tree) の定義

class BIT:
  def __init__(self, n): pass
  def sum(self, i): pass
  def add(self, i, x) pass
  def lowerbound(self, s): pass

sum(i): index $0$ から $i$ までの和を返します
add(i, x): index $i$ に $x$ を足します
lowerbound(s): index $0$ から $i$ までの和が $s$ になるような最小の $i$ を返します

いずれの操作も、列の長さ $N$ に対して時間計算量 $\mathcal{O}(\log{N})$ で実現できることが知られています。

BIT を拡張して BITSortedMultiSet を作る

class BITSortedMultiSet:
  def __init__(self, mx):
    self.mx = mx
    self.fwk = BIT(mx)
  def add(self, x): self.fwk.add(x, 1)
  def discard(self, x):
    if self.fwk.sum(x) - self.fwk.sum(x - 1): self.fwk.add(x, -1)
  def kthmax(self, k):
    n = self.fwk.sum(self.mx)
    return self.fwk.lowerbound(n - k + 1)

受け付ける最大の整数をコンストラクタに渡す必要があります。

add(x): 集合に $x$ を追加します
discard(x): 集合から $x$ をひとつ減らします
kthmax(k): $k$ 番目に大きい値を返します

制約

BIT の制約から $1$ 未満を受け付けることはできません。また、整数以外受け付けられません。

おわり

よく使います。問題に特化したデータ構造の拡張をしておくと、脳のリソースをセーブできて良いです。