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>
This commit is contained in:
Saksham Chawla 2022-10-12 22:05:31 +05:30 committed by GitHub
parent a04a6365de
commit d15bf7d492
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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 @@ class Heap:
# 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 @@ class Heap:
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 @@ class Heap:
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 @@ class Heap:
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 @@ class Heap:
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