heap.py

from __future__ importannotations

fromabcimportabstractmethod
fromcollections.abcimportIterable
fromtypingimportGeneric, Protocol, TypeVar


classComparable(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)


classHeap(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:
returnstr(self.h)

defparent_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'
 """
ifchild_idx>0:
return (child_idx-1) //2
returnNone

defleft_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
ifleft_child_index<self.heap_size:
returnleft_child_index
returnNone

defright_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
ifright_child_index<self.heap_size:
returnright_child_index
returnNone

defmax_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.
 """
ifindex<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
ifleft_childisnotNoneandself.h[left_child] >self.h[violation]:
violation=left_child
ifright_childisnotNoneandself.h[right_child] >self.h[violation]:
violation=right_child
# if violation indeed exists
ifviolation!=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)

defbuild_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)
ifself.heap_size>1:
# max_heapify from right to left but exclude leaves (last level)
foriinrange(self.heap_size//2-1, -1, -1):
self.max_heapify(i)

defextract_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
 """
ifself.heap_size>=2:
me=self.h[0]
self.h[0] =self.h.pop(-1)
self.heap_size-=1
self.max_heapify(0)
returnme
elifself.heap_size==1:
self.heap_size-=1
returnself.h.pop(-1)
else:
raiseException("Empty heap")

definsert(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
whileidx>=0:
self.max_heapify(idx)
idx= (idx-1) //2

defheap_sort(self) ->None:
size=self.heap_size
forjinrange(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__":
importdoctest

# run doc test
doctest.testmod()

# demo
forunsortedin [
 [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")