ErasableHeap

ErasableHeap.py

from heapq import heappush, heappop, heapify

class ErasableHeap(object):
    # https://qiita.com/Salmonize/items/638da118cd621d2628d1

    def __init__(self, minheap=True):
        self.data = []
        self.event = []

        if minheap:
            self.c = 1              # 最小ヒープ
        else:
            self.c = -1

    def push(self, x):
        return heappush(self.data, self.c * x)

    def erase(self, x):
        return heappush(self.event, self.c * x)

    def _check(self):
        while self.event and\
                self.data[0] == self.event[0]:
            heappop(self.data)
            heappop(self.event)

    def front(self):
        self._check()
        if len(self.data) == 0:
            return None
        return self.c * self.data[0]

    def pop(self):
        self.check()
        return self.c * heappop(self.data)

    def __len__(self):
        return len(self.data)