削除可能ヒープ

heapq のヒープとは別に「削除を待っている値」のヒープを持ち、最小値を見るときに両方の先頭が一致する間だけ取り除きます。同じ値が複数入っていても使えます。

• ErasableHeap(a): 列 $a$ で初期化（省略すると空）
• push(x): $x$ を追加
• erase(x): $x$ を 1 つ削除（ $x$ が入っているときだけ使う）
• top(): 最小値
• pop(): 最小値を削除して返す
• len(h): 要素数
• 計算量: 各操作ならし $O(\log N)$

最大値を扱いたいときは、符号を反転した値を入れます。

コード

from heapq import heapify, heappush, heappop


class ErasableHeap:
    def __init__(self, a=()):
        self.hq = list(a)
        heapify(self.hq)
        self.dq = []  # 削除を待っている値
        self.size = len(self.hq)

    def push(self, x):
        heappush(self.hq, x)
        self.size += 1

    def erase(self, x):
        # x が入っているときだけ使う
        heappush(self.dq, x)
        self.size -= 1

    def _clean(self):
        while self.dq and self.hq[0] == self.dq[0]:
            heappop(self.hq)
            heappop(self.dq)

    def top(self):
        self._clean()
        return self.hq[0]

    def pop(self):
        self._clean()
        self.size -= 1
        return heappop(self.hq)

    def __len__(self):
        return self.size