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