Add typing to data_structures/heap/heap_generic.py

data_structures/heap/heap_generic.py

@@ -1,35 +1,38 @@
+from collections.abc import Callable
+
+
 class Heap:
     """
     A generic Heap class, can be used as min or max by passing the key function
     accordingly.
     """
 
-    def __init__(self, key=None):
+    def __init__(self, key: Callable | None = None) -> None:
         # Stores actual heap items.
-        self.arr = list()
+        self.arr: list = list()
         # Stores indexes of each item for supporting updates and deletion.
-        self.pos_map = {}
+        self.pos_map: dict = {}
         # Stores current size of heap.
         self.size = 0
         # Stores function used to evaluate the score of an item on which basis ordering
         # will be done.
         self.key = key or (lambda x: x)
 
-    def _parent(self, i):
+    def _parent(self, i: int) -> int | None:
         """Returns parent index of given index if exists else None"""
         return int((i - 1) / 2) if i > 0 else None
 
-    def _left(self, i):
+    def _left(self, i: int) -> int | None:
         """Returns left-child-index of given index if exists else None"""
         left = int(2 * i + 1)
         return left if 0 < left < self.size else None
 
-    def _right(self, i):
+    def _right(self, i: int) -> int | None:
         """Returns right-child-index of given index if exists else None"""
         right = int(2 * i + 2)
         return right if 0 < right < self.size else None
 
-    def _swap(self, i, j):
+    def _swap(self, i: int, j: int) -> None:
         """Performs changes required for swapping two elements in the heap"""
         # First update the indexes of the items in index map.
         self.pos_map[self.arr[i][0]], self.pos_map[self.arr[j][0]] = (
@@ -39,11 +42,11 @@ def _swap(self, i, j):
         # Then swap the items in the list.
         self.arr[i], self.arr[j] = self.arr[j], self.arr[i]
 
-    def _cmp(self, i, j):
+    def _cmp(self, i: int, j: int) -> bool:
         """Compares the two items using default comparison"""
         return self.arr[i][1] < self.arr[j][1]
 
-    def _get_valid_parent(self, i):
+    def _get_valid_parent(self, i: int) -> int:
         """
         Returns index of valid parent as per desired ordering among given index and
         both it's children
@@ -59,21 +62,21 @@ def _get_valid_parent(self, i):
 
         return valid_parent
 
-    def _heapify_up(self, index):
+    def _heapify_up(self, index: int) -> None:
         """Fixes the heap in upward direction of given index"""
         parent = self._parent(index)
         while parent is not None and not self._cmp(index, parent):
             self._swap(index, parent)
             index, parent = parent, self._parent(parent)
 
-    def _heapify_down(self, index):
+    def _heapify_down(self, index: int) -> None:
         """Fixes the heap in downward direction of given index"""
         valid_parent = self._get_valid_parent(index)
         while valid_parent != index:
             self._swap(index, valid_parent)
             index, valid_parent = valid_parent, self._get_valid_parent(valid_parent)
 
-    def update_item(self, item, item_value):
+    def update_item(self, item: int, item_value: int) -> None:
         """Updates given item value in heap if present"""
         if item not in self.pos_map:
             return
@@ -84,7 +87,7 @@ def update_item(self, item, item_value):
         self._heapify_up(index)
         self._heapify_down(index)
 
-    def delete_item(self, item):
+    def delete_item(self, item: int) -> None:
         """Deletes given item from heap if present"""
         if item not in self.pos_map:
             return
@@ -99,7 +102,7 @@ def delete_item(self, item):
             self._heapify_up(index)
             self._heapify_down(index)
 
-    def insert_item(self, item, item_value):
+    def insert_item(self, item: int, item_value: int) -> None:
         """Inserts given item with given value in heap"""
         arr_len = len(self.arr)
         if arr_len == self.size:
@@ -110,11 +113,11 @@ def insert_item(self, item, item_value):
         self.size += 1
         self._heapify_up(self.size - 1)
 
-    def get_top(self):
+    def get_top(self) -> tuple | None:
         """Returns top item tuple (Calculated value, item) from heap if present"""
         return self.arr[0] if self.size else None
 
-    def extract_top(self):
+    def extract_top(self) -> tuple | None:
         """
         Return top item tuple (Calculated value, item) from heap and removes it as well
         if present