Source code for titan_pylib.data_structures.heap.deletable_min_heap

 1from titan_pylib.data_structures.heap.min_heap import MinHeap
 2from titan_pylib.my_class.supports_less_than import SupportsLessThan
 3from typing import TypeVar, Generic, Iterable
 4
 5T = TypeVar("T", bound=SupportsLessThan)
 6
 7

[docs]
 8class DeletableMinHeap(Generic[T]):
 9
10    def __init__(self, a: Iterable[T] = []) -> None:
11        """削除可能Minヒープです。
12        要素の 追加/削除/最小値取得 が効率よく行えます。
13        """
14        self.hq: MinHeap[T] = MinHeap(a)
15        self.rem_hq: MinHeap[T] = MinHeap()
16        self._len: int = len(self.hq)
17

[docs]
18    def push(self, key: T) -> None:
19        """``key`` を追加します。
20        :math:`O(\\log{n})` です。
21        """
22        self._len += 1
23        if self.rem_hq and self.rem_hq.get_min() == key:
24            self.rem_hq.pop_min()
25            return
26        self.hq.push(key)

27

[docs]
28    def remove(self, key: T) -> None:
29        """``key`` を削除します。
30        :math:`O(\\log{n})` です。
31        """
32        assert self._len > 0
33        self._len -= 1
34        if self.hq.get_min() == key:
35            self.hq.pop_min()
36        else:
37            self.rem_hq.push(key)

38
39    def _delete(self) -> None:
40        while self.rem_hq and self.rem_hq.get_min() == self.hq.get_min():
41            self.hq.pop_min()
42            self.rem_hq.pop_min()
43

[docs]
44    def get_min(self) -> T:
45        """最小の要素を返します。
46        :math:`O(\\log{n})` です。
47        """
48        assert self._len > 0
49        self._delete()
50        return self.hq.get_min()

51

[docs]
52    def pop_min(self) -> T:
53        """最小の要素を削除して返します。
54        :math:`O(\\log{n})` です。
55        """
56        assert self._len > 0
57        self._len -= 1
58        self._delete()
59        return self.hq.pop_min()

60
61    def __len__(self) -> int:
62        return self._len