1
+ from collections .abc import Callable
2
+
3
+
1
4
class Heap :
2
5
"""
3
6
A generic Heap class, can be used as min or max by passing the key function
4
7
accordingly.
5
8
"""
6
9
7
- def __init__ (self , key = None ):
10
+ def __init__ (self , key : Callable | None = None ) -> None :
8
11
# Stores actual heap items.
9
- self .arr = list ()
12
+ self .arr : list = list ()
10
13
# Stores indexes of each item for supporting updates and deletion.
11
- self .pos_map = {}
14
+ self .pos_map : dict = {}
12
15
# Stores current size of heap.
13
16
self .size = 0
14
17
# Stores function used to evaluate the score of an item on which basis ordering
15
18
# will be done.
16
19
self .key = key or (lambda x : x )
17
20
18
- def _parent (self , i ) :
21
+ def _parent (self , i : int ) -> int | None :
19
22
"""Returns parent index of given index if exists else None"""
20
23
return int ((i - 1 ) / 2 ) if i > 0 else None
21
24
22
- def _left (self , i ) :
25
+ def _left (self , i : int ) -> int | None :
23
26
"""Returns left-child-index of given index if exists else None"""
24
27
left = int (2 * i + 1 )
25
28
return left if 0 < left < self .size else None
26
29
27
- def _right (self , i ) :
30
+ def _right (self , i : int ) -> int | None :
28
31
"""Returns right-child-index of given index if exists else None"""
29
32
right = int (2 * i + 2 )
30
33
return right if 0 < right < self .size else None
31
34
32
- def _swap (self , i , j ) :
35
+ def _swap (self , i : int , j : int ) -> None :
33
36
"""Performs changes required for swapping two elements in the heap"""
34
37
# First update the indexes of the items in index map.
35
38
self .pos_map [self .arr [i ][0 ]], self .pos_map [self .arr [j ][0 ]] = (
@@ -39,11 +42,11 @@ def _swap(self, i, j):
39
42
# Then swap the items in the list.
40
43
self .arr [i ], self .arr [j ] = self .arr [j ], self .arr [i ]
41
44
42
- def _cmp (self , i , j ) :
45
+ def _cmp (self , i : int , j : int ) -> bool :
43
46
"""Compares the two items using default comparison"""
44
47
return self .arr [i ][1 ] < self .arr [j ][1 ]
45
48
46
- def _get_valid_parent (self , i ) :
49
+ def _get_valid_parent (self , i : int ) -> int :
47
50
"""
48
51
Returns index of valid parent as per desired ordering among given index and
49
52
both it's children
@@ -59,21 +62,21 @@ def _get_valid_parent(self, i):
59
62
60
63
return valid_parent
61
64
62
- def _heapify_up (self , index ) :
65
+ def _heapify_up (self , index : int ) -> None :
63
66
"""Fixes the heap in upward direction of given index"""
64
67
parent = self ._parent (index )
65
68
while parent is not None and not self ._cmp (index , parent ):
66
69
self ._swap (index , parent )
67
70
index , parent = parent , self ._parent (parent )
68
71
69
- def _heapify_down (self , index ) :
72
+ def _heapify_down (self , index : int ) -> None :
70
73
"""Fixes the heap in downward direction of given index"""
71
74
valid_parent = self ._get_valid_parent (index )
72
75
while valid_parent != index :
73
76
self ._swap (index , valid_parent )
74
77
index , valid_parent = valid_parent , self ._get_valid_parent (valid_parent )
75
78
76
- def update_item (self , item , item_value ) :
79
+ def update_item (self , item : int , item_value : int ) -> None :
77
80
"""Updates given item value in heap if present"""
78
81
if item not in self .pos_map :
79
82
return
@@ -84,7 +87,7 @@ def update_item(self, item, item_value):
84
87
self ._heapify_up (index )
85
88
self ._heapify_down (index )
86
89
87
- def delete_item (self , item ) :
90
+ def delete_item (self , item : int ) -> None :
88
91
"""Deletes given item from heap if present"""
89
92
if item not in self .pos_map :
90
93
return
@@ -99,7 +102,7 @@ def delete_item(self, item):
99
102
self ._heapify_up (index )
100
103
self ._heapify_down (index )
101
104
102
- def insert_item (self , item , item_value ) :
105
+ def insert_item (self , item : int , item_value : int ) -> None :
103
106
"""Inserts given item with given value in heap"""
104
107
arr_len = len (self .arr )
105
108
if arr_len == self .size :
@@ -110,11 +113,11 @@ def insert_item(self, item, item_value):
110
113
self .size += 1
111
114
self ._heapify_up (self .size - 1 )
112
115
113
- def get_top (self ):
116
+ def get_top (self ) -> tuple | None :
114
117
"""Returns top item tuple (Calculated value, item) from heap if present"""
115
118
return self .arr [0 ] if self .size else None
116
119
117
- def extract_top (self ):
120
+ def extract_top (self ) -> tuple | None :
118
121
"""
119
122
Return top item tuple (Calculated value, item) from heap and removes it as well
120
123
if present
0 commit comments