Source code for titan_pylib.data_structures.heap.deletable_max_heap

 1from titan_pylib.data_structures.heap.max_heap import MaxHeap
 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 DeletableMaxHeap(Generic[T]):
 9
10    def __init__(self, a: Iterable[T] = []) -> None:
11        """削除可能Maxヒープです。
12        要素の 追加/削除/最大値取得 が効率よく行えます。
13        """
14        self.hq: MaxHeap[T] = MaxHeap(a)
15        self.rem_hq: MaxHeap[T] = MaxHeap()
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_max() == key:
24            self.rem_hq.pop_max()
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_max() == key:
35            self.hq.pop_max()
36        else:
37            self.rem_hq.push(key)

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

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

51

[docs]
52    def pop_max(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_max()

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