Skip to content

Commit d15bf7d

Browse files
Add typing to data_structures/heap/heap_generic.py (#7044)
* Update heap_generic.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update heap_generic.py * Update heap_generic.py * Update heap_generic.py * Update heap_generic.py * Update heap_generic.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent a04a636 commit d15bf7d

File tree

1 file changed

+19
-16
lines changed

1 file changed

+19
-16
lines changed

data_structures/heap/heap_generic.py

+19-16
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,38 @@
1+
from collections.abc import Callable
2+
3+
14
class Heap:
25
"""
36
A generic Heap class, can be used as min or max by passing the key function
47
accordingly.
58
"""
69

7-
def __init__(self, key=None):
10+
def __init__(self, key: Callable | None = None) -> None:
811
# Stores actual heap items.
9-
self.arr = list()
12+
self.arr: list = list()
1013
# Stores indexes of each item for supporting updates and deletion.
11-
self.pos_map = {}
14+
self.pos_map: dict = {}
1215
# Stores current size of heap.
1316
self.size = 0
1417
# Stores function used to evaluate the score of an item on which basis ordering
1518
# will be done.
1619
self.key = key or (lambda x: x)
1720

18-
def _parent(self, i):
21+
def _parent(self, i: int) -> int | None:
1922
"""Returns parent index of given index if exists else None"""
2023
return int((i - 1) / 2) if i > 0 else None
2124

22-
def _left(self, i):
25+
def _left(self, i: int) -> int | None:
2326
"""Returns left-child-index of given index if exists else None"""
2427
left = int(2 * i + 1)
2528
return left if 0 < left < self.size else None
2629

27-
def _right(self, i):
30+
def _right(self, i: int) -> int | None:
2831
"""Returns right-child-index of given index if exists else None"""
2932
right = int(2 * i + 2)
3033
return right if 0 < right < self.size else None
3134

32-
def _swap(self, i, j):
35+
def _swap(self, i: int, j: int) -> None:
3336
"""Performs changes required for swapping two elements in the heap"""
3437
# First update the indexes of the items in index map.
3538
self.pos_map[self.arr[i][0]], self.pos_map[self.arr[j][0]] = (
@@ -39,11 +42,11 @@ def _swap(self, i, j):
3942
# Then swap the items in the list.
4043
self.arr[i], self.arr[j] = self.arr[j], self.arr[i]
4144

42-
def _cmp(self, i, j):
45+
def _cmp(self, i: int, j: int) -> bool:
4346
"""Compares the two items using default comparison"""
4447
return self.arr[i][1] < self.arr[j][1]
4548

46-
def _get_valid_parent(self, i):
49+
def _get_valid_parent(self, i: int) -> int:
4750
"""
4851
Returns index of valid parent as per desired ordering among given index and
4952
both it's children
@@ -59,21 +62,21 @@ def _get_valid_parent(self, i):
5962

6063
return valid_parent
6164

62-
def _heapify_up(self, index):
65+
def _heapify_up(self, index: int) -> None:
6366
"""Fixes the heap in upward direction of given index"""
6467
parent = self._parent(index)
6568
while parent is not None and not self._cmp(index, parent):
6669
self._swap(index, parent)
6770
index, parent = parent, self._parent(parent)
6871

69-
def _heapify_down(self, index):
72+
def _heapify_down(self, index: int) -> None:
7073
"""Fixes the heap in downward direction of given index"""
7174
valid_parent = self._get_valid_parent(index)
7275
while valid_parent != index:
7376
self._swap(index, valid_parent)
7477
index, valid_parent = valid_parent, self._get_valid_parent(valid_parent)
7578

76-
def update_item(self, item, item_value):
79+
def update_item(self, item: int, item_value: int) -> None:
7780
"""Updates given item value in heap if present"""
7881
if item not in self.pos_map:
7982
return
@@ -84,7 +87,7 @@ def update_item(self, item, item_value):
8487
self._heapify_up(index)
8588
self._heapify_down(index)
8689

87-
def delete_item(self, item):
90+
def delete_item(self, item: int) -> None:
8891
"""Deletes given item from heap if present"""
8992
if item not in self.pos_map:
9093
return
@@ -99,7 +102,7 @@ def delete_item(self, item):
99102
self._heapify_up(index)
100103
self._heapify_down(index)
101104

102-
def insert_item(self, item, item_value):
105+
def insert_item(self, item: int, item_value: int) -> None:
103106
"""Inserts given item with given value in heap"""
104107
arr_len = len(self.arr)
105108
if arr_len == self.size:
@@ -110,11 +113,11 @@ def insert_item(self, item, item_value):
110113
self.size += 1
111114
self._heapify_up(self.size - 1)
112115

113-
def get_top(self):
116+
def get_top(self) -> tuple | None:
114117
"""Returns top item tuple (Calculated value, item) from heap if present"""
115118
return self.arr[0] if self.size else None
116119

117-
def extract_top(self):
120+
def extract_top(self) -> tuple | None:
118121
"""
119122
Return top item tuple (Calculated value, item) from heap and removes it as well
120123
if present

0 commit comments

Comments
 (0)