1- def _left (i ):
1+ """
2+ This module implements a max-heap data structure. It provides methods to add
3+ elements, pop the maximum element, and retrieve the top element without
4+ removing it. It is designed to be used in scenarios where you need to maintain
5+ a collection of elements and frequently access the maximum element
6+ efficiently.
7+ """
8+
9+ from __future__ import annotations
10+
11+ from typing import Any , Generic , Iterator , Optional , Protocol , TypeVar
12+
13+
14+ class Comparable (Protocol ):
15+ """
16+ A protocol that defines the methods required for an object to be
17+ orderable. This is used to ensure that the keys in the heap can be
18+ compared with each other.
19+ """
20+
21+ def __lt__ (self : K , other : K ) -> bool : ...
22+
23+ def __le__ (self : K , other : K ) -> bool : ...
24+
25+ def __gt__ (self : K , other : K ) -> bool : ...
26+
27+ def __ge__ (self : K , other : K ) -> bool : ...
28+
29+
30+ K = TypeVar ("K" , bound = Comparable )
31+
32+ V = TypeVar ("V" )
33+
34+
35+ def _left (i : int ) -> int :
36+ """
37+ Returns the index of the left child of the node at index i in a binary heap.
38+
39+ :param i: The index of the parent node.
40+ :return: The index of the left child.
41+ """
242 return 2 * i + 1
343
444
5- def _right (i ):
45+ def _right (i : int ) -> int :
46+ """
47+ Returns the index of the right child of the node at index i in a binary heap.
48+
49+ :param i: The index of the parent node.
50+ :return: The index of the right child.
51+ """
652 return 2 * i + 2
753
854
9- def _parent (i ):
55+ def _parent (i : int ) -> int :
56+ """
57+ Returns the index of the parent node of the node at index i in a binary heap.
58+
59+ :param i: The index of the child node.
60+ :return: The index of the parent node.
61+ """
1062 return max (0 , (i - 1 ) // 2 )
1163
1264
13- class _HeapNode :
65+ class _HeapNode (Generic [K , V ]):
66+ """
67+ A private class representing a node in the max-heap. Each node contains a
68+ key and a value. The key is used to determine the order of the nodes in
69+ the heap, with larger keys being prioritized over smaller keys. The value
70+ can be any associated data that you want to store with the key.
71+
72+ :param key: The key of the node, used for ordering in the heap.
73+ :param value: The value associated with the key.
74+ """
1475
15- def __init__ (self , key , value ) :
76+ def __init__ (self , key : K , value : V ) -> None :
1677 self ._key = key
1778 self ._value = value
1879
19- def get (self ):
80+ def get (self ) -> tuple [K , V ]:
81+ """
82+ Returns the key and value of the node as a tuple.
83+
84+ :return: A tuple containing the key and value of the node.
85+ """
2086 return self ._key , self ._value
2187
22- def __lt__ (self , other ) :
88+ def __lt__ (self , other : _HeapNode [ K , Any ]) -> bool :
2389 return self ._key < other ._key
2490
25- def __le__ (self , other ) :
91+ def __le__ (self , other : _HeapNode [ K , Any ]) -> bool :
2692 return self ._key <= other ._key
2793
28- def __gt__ (self , other ) :
94+ def __gt__ (self , other : _HeapNode [ K , Any ]) -> bool :
2995 return self ._key > other ._key
3096
31- def __ge__ (self , other ) :
97+ def __ge__ (self , other : _HeapNode [ K , Any ]) -> bool :
3298 return self ._key >= other ._key
3399
34100
35- class MaxHeap :
101+ class MaxHeap (Generic [K , V ]):
102+ """
103+ A max-heap implementation that allows for efficient retrieval and removal
104+ of the maximum element. The heap is implemented as a list of _HeapNode
105+ objects, where each node contains a key and a value. The key is used to
106+ determine the order of the elements in the heap, with larger keys being
107+ prioritized over smaller keys.
108+ """
36109
37- def __init__ (self ):
110+ _heap : list [_HeapNode [K , V ]]
111+ _iter : Iterator [_HeapNode [K , V ]]
112+
113+ def __init__ (self ) -> None :
38114 self ._heap = []
39- self ._iter = None
40115
41- def __iter__ (self ):
116+ def __iter__ (self ) -> MaxHeap [ K , V ] :
42117 self ._iter = iter (self ._heap )
43118 return self
44119
45- def __next__ (self ):
120+ def __next__ (self ) -> tuple [ K , V ] :
46121 node = next (self ._iter )
47122 return node .get ()
48123
49- def __len__ (self ):
124+ def __len__ (self ) -> int :
50125 return len (self ._heap )
51126
52- def top (self ):
127+ def top (self ) -> tuple [Optional [K ], Optional [V ]]:
128+ """
129+ Returns the maximum element of the heap without removing it.
130+
131+ :return: A tuple containing the key and value of the maximum element,
132+ or (None, None) if the heap is empty.
133+ """
53134 if not self ._heap :
54135 return (None , None )
55136 return self ._heap [0 ].get ()
56137
57- def pop (self ):
138+ def pop (self ) -> tuple [Optional [K ], Optional [V ]]:
139+ """
140+ Removes and returns the maximum element from the heap.
141+
142+ :return: A tuple containing the key and value of the maximum element,
143+ or (None, None) if the heap is empty.
144+ """
58145 if not self ._heap :
59- return
146+ return ( None , None )
60147 max_val = self ._heap [0 ]
61148 self ._heap [0 ] = self ._heap [- 1 ]
62149 self ._heap .pop ()
63150 self ._bubble_down ()
64151 return max_val .get ()
65152
66- def add (self , key , val ):
153+ def add (self , key : K , val : V ) -> None :
154+ """
155+ Adds a new element to the heap.
156+
157+ :param key: The key of the element to be added, which determines its
158+ position in the heap.
159+ :param val: The value associated with the key.
160+ :return: None
161+ """
67162 self ._heap .append (_HeapNode (key , val ))
68163 self ._bubble_up ()
69164
70- def _get_local_max (self , i ) :
165+ def _get_local_max (self , i : int ) -> int :
71166 heap_len = len (self ._heap )
72167 left = _left (i )
73168 right = _right (i )
@@ -84,7 +179,7 @@ def _get_local_max(self, i):
84179 return max_child
85180 return i
86181
87- def _fix_down (self , i ) :
182+ def _fix_down (self , i : int ) -> int :
88183 local_max = self ._get_local_max (i )
89184 if i < local_max :
90185 self ._heap [i ], self ._heap [local_max ] = (
@@ -94,22 +189,22 @@ def _fix_down(self, i):
94189 return local_max
95190 return i
96191
97- def _fix_up (self , i ) :
192+ def _fix_up (self , i : int ) -> int :
98193 parent = _parent (i )
99194 if self ._heap [parent ] < self ._heap [i ]:
100195 self ._heap [i ], self ._heap [parent ] = self ._heap [parent ], self ._heap [i ]
101196 return parent
102197 return i
103198
104- def _bubble_down (self ):
199+ def _bubble_down (self ) -> None :
105200 current = 0
106201 done = False
107202 while not done :
108203 local_max = self ._fix_down (current )
109204 done = current == local_max
110205 current = local_max
111206
112- def _bubble_up (self ):
207+ def _bubble_up (self ) -> None :
113208 current = len (self ._heap ) - 1
114209 done = False
115210 while not done :
0 commit comments