from __future__ import annotations from abc import abstractmethod from collections.abc import Iterable from typing import Generic, Protocol, TypeVar class Comparable(Protocol): @abstractmethod def __lt__(self: T, other: T) -> bool: pass @abstractmethod def __gt__(self: T, other: T) -> bool: pass @abstractmethod def __eq__(self: T, other: object) -> bool: pass T = TypeVar("T", bound=Comparable) class Heap(Generic[T]): """A Max Heap Implementation >>> unsorted = [103, 9, 1, 7, 11, 15, 25, 201, 209, 107, 5] >>> h = Heap() >>> h.build_max_heap(unsorted) >>> h [209, 201, 25, 103, 107, 15, 1, 9, 7, 11, 5] >>> >>> h.extract_max() 209 >>> h [201, 107, 25, 103, 11, 15, 1, 9, 7, 5] >>> >>> h.insert(100) >>> h [201, 107, 25, 103, 100, 15, 1, 9, 7, 5, 11] >>> >>> h.heap_sort() >>> h [1, 5, 7, 9, 11, 15, 25, 100, 103, 107, 201] """ def __init__(self) -> None: self.h: list[T] = [] self.heap_size: int = 0 def __repr__(self) -> str: return str(self.h) def parent_index(self, child_idx: int) -> int | None: """ returns the parent index based on the given child index >>> h = Heap() >>> h.build_max_heap([103, 9, 1, 7, 11, 15, 25, 201, 209, 107, 5]) >>> h [209, 201, 25, 103, 107, 15, 1, 9, 7, 11, 5] >>> h.parent_index(-1) # returns none if index is <=0 >>> h.parent_index(0) # returns none if index is <=0 >>> h.parent_index(1) 0 >>> h.parent_index(2) 0 >>> h.parent_index(3) 1 >>> h.parent_index(4) 1 >>> h.parent_index(5) 2 >>> h.parent_index(10.5) 4.0 >>> h.parent_index(209.0) 104.0 >>> h.parent_index("Test") Traceback (most recent call last): ... TypeError: '>' not supported between instances of 'str' and 'int' """ if child_idx > 0: return (child_idx - 1) // 2 return None def left_child_idx(self, parent_idx: int) -> int | None: """ return the left child index if the left child exists. if not, return None. """ left_child_index = 2 * parent_idx + 1 if left_child_index < self.heap_size: return left_child_index return None def right_child_idx(self, parent_idx: int) -> int | None: """ return the right child index if the right child exists. if not, return None. """ right_child_index = 2 * parent_idx + 2 if right_child_index < self.heap_size: return right_child_index return None def max_heapify(self, index: int) -> None: """ correct a single violation of the heap property in a subtree's root. It is the function that is responsible for restoring the property of Max heap i.e the maximum element is always at top. """ if index < self.heap_size: violation: int = index left_child = self.left_child_idx(index) right_child = self.right_child_idx(index) # check which child is larger than its parent if left_child is not None and self.h[left_child] > self.h[violation]: violation = left_child if right_child is not None and self.h[right_child] > self.h[violation]: violation = right_child # if violation indeed exists if violation != index: # swap to fix the violation self.h[violation], self.h[index] = self.h[index], self.h[violation] # fix the subsequent violation recursively if any self.max_heapify(violation) def build_max_heap(self, collection: Iterable[T]) -> None: """ build max heap from an unsorted array >>> h = Heap() >>> h.build_max_heap([20,40,50,20,10]) >>> h [50, 40, 20, 20, 10] >>> h = Heap() >>> h.build_max_heap([1,2,3,4,5,6,7,8,9,0]) >>> h [9, 8, 7, 4, 5, 6, 3, 2, 1, 0] >>> h = Heap() >>> h.build_max_heap([514,5,61,57,8,99,105]) >>> h [514, 57, 105, 5, 8, 99, 61] >>> h = Heap() >>> h.build_max_heap([514,5,61.6,57,8,9.9,105]) >>> h [514, 57, 105, 5, 8, 9.9, 61.6] """ self.h = list(collection) self.heap_size = len(self.h) if self.heap_size > 1: # max_heapify from right to left but exclude leaves (last level) for i in range(self.heap_size // 2 - 1, -1, -1): self.max_heapify(i) def extract_max(self) -> T: """ get and remove max from heap >>> h = Heap() >>> h.build_max_heap([20,40,50,20,10]) >>> h.extract_max() 50 >>> h = Heap() >>> h.build_max_heap([514,5,61,57,8,99,105]) >>> h.extract_max() 514 >>> h = Heap() >>> h.build_max_heap([1,2,3,4,5,6,7,8,9,0]) >>> h.extract_max() 9 """ if self.heap_size >= 2: me = self.h[0] self.h[0] = self.h.pop(-1) self.heap_size -= 1 self.max_heapify(0) return me elif self.heap_size == 1: self.heap_size -= 1 return self.h.pop(-1) else: raise Exception("Empty heap") def insert(self, value: T) -> None: """ insert a new value into the max heap >>> h = Heap() >>> h.insert(10) >>> h [10] >>> h = Heap() >>> h.insert(10) >>> h.insert(10) >>> h [10, 10] >>> h = Heap() >>> h.insert(10) >>> h.insert(10.1) >>> h [10.1, 10] >>> h = Heap() >>> h.insert(0.1) >>> h.insert(0) >>> h.insert(9) >>> h.insert(5) >>> h [9, 5, 0.1, 0] """ self.h.append(value) idx = (self.heap_size - 1) // 2 self.heap_size += 1 while idx >= 0: self.max_heapify(idx) idx = (idx - 1) // 2 def heap_sort(self) -> None: size = self.heap_size for j in range(size - 1, 0, -1): self.h[0], self.h[j] = self.h[j], self.h[0] self.heap_size -= 1 self.max_heapify(0) self.heap_size = size if __name__ == "__main__": import doctest # run doc test doctest.testmod() # demo for unsorted in [ [0], [2], [3, 5], [5, 3], [5, 5], [0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 3, 5], [0, 2, 2, 3, 5], [2, 5, 3, 0, 2, 3, 0, 3], [6, 1, 2, 7, 9, 3, 4, 5, 10, 8], [103, 9, 1, 7, 11, 15, 25, 201, 209, 107, 5], [-45, -2, -5], ]: print(f"unsorted array: {unsorted}") heap: Heap[int] = Heap() heap.build_max_heap(unsorted) print(f"after build heap: {heap}") print(f"max value: {heap.extract_max()}") print(f"after max value removed: {heap}") heap.insert(100) print(f"after new value 100 inserted: {heap}") heap.heap_sort() print(f"heap-sorted array: {heap}\n")