Added doctest to heap.py (#11059)

This commit is contained in:
Suyash Dongre 2023-10-29 01:13:20 +05:30 committed by GitHub
parent f2436318ce
commit b51b833e0a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -81,6 +81,9 @@ class Heap(Generic[T]):
def max_heapify(self, index: int) -> None: def max_heapify(self, index: int) -> None:
""" """
correct a single violation of the heap property in a subtree's root. 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: if index < self.heap_size:
violation: int = index violation: int = index
@ -99,7 +102,29 @@ class Heap(Generic[T]):
self.max_heapify(violation) self.max_heapify(violation)
def build_max_heap(self, collection: Iterable[T]) -> None: def build_max_heap(self, collection: Iterable[T]) -> None:
"""build max heap from an unsorted array""" """
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.h = list(collection)
self.heap_size = len(self.h) self.heap_size = len(self.h)
if self.heap_size > 1: if self.heap_size > 1:
@ -108,7 +133,24 @@ class Heap(Generic[T]):
self.max_heapify(i) self.max_heapify(i)
def extract_max(self) -> T: def extract_max(self) -> T:
"""get and remove max from heap""" """
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: if self.heap_size >= 2:
me = self.h[0] me = self.h[0]
self.h[0] = self.h.pop(-1) self.h[0] = self.h.pop(-1)
@ -122,7 +164,34 @@ class Heap(Generic[T]):
raise Exception("Empty heap") raise Exception("Empty heap")
def insert(self, value: T) -> None: def insert(self, value: T) -> None:
"""insert a new value into the max heap""" """
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) self.h.append(value)
idx = (self.heap_size - 1) // 2 idx = (self.heap_size - 1) // 2
self.heap_size += 1 self.heap_size += 1