Source code for titan_pylib.data_structures.heap.max_heap

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

[docs]
 7class MaxHeap(Generic[T]):
 8
 9    def __init__(self, a: Iterable[T] = []) -> None:
10        self.a = list(a)
11        self._heapify()
12
13    def _heapify(self) -> None:
14        for i in range(len(self.a) - 1, -1, -1):
15            self._down(i)
16
17    def _down(self, i: int) -> None:
18        a = self.a
19        n = len(a)
20        while i * 2 + 1 < n:
21            u, v = i * 2 + 1, i * 2 + 2
22            if v < n and a[u] < a[v]:
23                u = v
24            if a[i] < a[u]:
25                a[i], a[u] = a[u], a[i]
26                i = u
27            else:
28                break
29
30    def _up(self, i: int) -> None:
31        a = self.a
32        while i > 0:
33            p = (i - 1) >> 1
34            if a[i] > a[p]:
35                a[i], a[p] = a[p], a[i]
36                i = p
37            else:
38                break
39

[docs]
40    def get_max(self) -> T:
41        return self.a[0]

42

[docs]
43    def pop_max(self) -> T:
44        res = self.a[0]
45        self.a[0] = self.a[-1]
46        self.a.pop()
47        self._down(0)
48        return res

49

[docs]
50    def push(self, key: T) -> None:
51        self.a.append(key)
52        self._up(len(self.a) - 1)

53

[docs]
54    def pushpop_max(self, key: T) -> T:
55        if self.a[0] < key or self.a[0] == key:
56            return key
57        res = self.a[0]
58        self.a[0] = key
59        self._down(0)
60        return res

61

[docs]
62    def replace_max(self, key: T) -> T:
63        res = self.a[0]
64        self.a[0] = key
65        self._down(0)
66        return res

67

[docs]
68    def tolist(self) -> list[T]:
69        return sorted(self.a)

70
71    def __getitem__(self, k: int) -> T:
72        assert k == 0
73        return self.a[0]
74
75    def __len__(self):
76        return len(self.a)
77
78    def __str__(self):
79        return str(self.a)
80
81    def __repr__(self):
82        return f"{self.__class__.__name__}({self})"