mirror of
https://github.com/TheAlgorithms/Python.git
synced 2024-12-25 12:40:14 +00:00
0e2e6abd6f
* Added doctest to heap.py * Update heap.py
277 lines
7.2 KiB
Python
277 lines
7.2 KiB
Python
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")
|