priorityqueue.py

import heapq
from collections.abc import Iterable
from typing import TypeVar

_T = TypeVar('_T')

# N UpdatableHeap
# I {"version": "1.1", "import": ["heapq"], "abc": ["Iterable"], "typing": ["TypeVar"], "const": ["_T"]}
_hpop, _hpush = heapq.heappop, heapq.heappush


class UpdatableHeap():
    """ Min-heap that supports delete and update operations."""

    __slots__ = ('_heap', '_del_heap')

    def __init__(self, values: Iterable[_T] = ()):
        self._heap = list(values)
        heapq.heapify(self._heap)
        self._del_heap = []

    def __len__(self):
        return len(self._heap) - len(self._del_heap)

    def push(self, val: _T):
        _hpush(self._heap, val)

    def pop(self) -> _T:
        p = _hpop(self._heap)
        while self._del_heap and self._del_heap[0] == self._heap[0]:
            _hpop(self._del_heap)
            _hpop(self._heap)
        return p

    def top(self) -> _T:
        return self._heap[0]

    def delete(self, val: _T):
        _hpush(self._del_heap, val)
        while self._del_heap and self._del_heap[0] == self._heap[0]:
            _hpop(self._del_heap)
            _hpop(self._heap)

    def update(self, val: _T, new_val: _T):
        _hpush(self._del_heap, val)
        _hpush(self._heap, new_val)

• 임의 원소 삭제를 지원하는 힙 참고