Skip to content

Commit 9823e26

Browse files
committed
Added doctest to heap.py
1 parent 1e1ee00 commit 9823e26

File tree

1 file changed

+72
-3
lines changed

1 file changed

+72
-3
lines changed

Diff for: data_structures/heap/heap.py

+72-3
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,9 @@ def right_child_idx(self, parent_idx: int) -> int | None:
8181
def max_heapify(self, index: int) -> None:
8282
"""
8383
correct a single violation of the heap property in a subtree's root.
84+
85+
It is the function that is responsible for restoring the property
86+
of Max heap i.e the maximum element is always at top.
8487
"""
8588
if index < self.heap_size:
8689
violation: int = index
@@ -99,7 +102,29 @@ def max_heapify(self, index: int) -> None:
99102
self.max_heapify(violation)
100103

101104
def build_max_heap(self, collection: Iterable[T]) -> None:
102-
"""build max heap from an unsorted array"""
105+
"""
106+
build max heap from an unsorted array
107+
108+
>>> h = Heap()
109+
>>> h.build_max_heap([20,40,50,20,10])
110+
>>> h
111+
[50, 40, 20, 20, 10]
112+
113+
>>> h = Heap()
114+
>>> h.build_max_heap([1,2,3,4,5,6,7,8,9,0])
115+
>>> h
116+
[9, 8, 7, 4, 5, 6, 3, 2, 1, 0]
117+
118+
>>> h = Heap()
119+
>>> h.build_max_heap([514,5,61,57,8,99,105])
120+
>>> h
121+
[514, 57, 105, 5, 8, 99, 61]
122+
123+
>>> h = Heap()
124+
>>> h.build_max_heap([514,5,61.6,57,8,9.9,105])
125+
>>> h
126+
[514, 57, 105, 5, 8, 9.9, 61.6]
127+
"""
103128
self.h = list(collection)
104129
self.heap_size = len(self.h)
105130
if self.heap_size > 1:
@@ -108,7 +133,24 @@ def build_max_heap(self, collection: Iterable[T]) -> None:
108133
self.max_heapify(i)
109134

110135
def extract_max(self) -> T:
111-
"""get and remove max from heap"""
136+
"""
137+
get and remove max from heap
138+
139+
>>> h = Heap()
140+
>>> h.build_max_heap([20,40,50,20,10])
141+
>>> h.extract_max()
142+
50
143+
144+
>>> h = Heap()
145+
>>> h.build_max_heap([514,5,61,57,8,99,105])
146+
>>> h.extract_max()
147+
514
148+
149+
>>> h = Heap()
150+
>>> h.build_max_heap([1,2,3,4,5,6,7,8,9,0])
151+
>>> h.extract_max()
152+
9
153+
"""
112154
if self.heap_size >= 2:
113155
me = self.h[0]
114156
self.h[0] = self.h.pop(-1)
@@ -122,7 +164,34 @@ def extract_max(self) -> T:
122164
raise Exception("Empty heap")
123165

124166
def insert(self, value: T) -> None:
125-
"""insert a new value into the max heap"""
167+
"""
168+
insert a new value into the max heap
169+
170+
>>> h = Heap()
171+
>>> h.insert(10)
172+
>>> h
173+
[10]
174+
175+
>>> h = Heap()
176+
>>> h.insert(10)
177+
>>> h.insert(10)
178+
>>> h
179+
[10, 10]
180+
181+
>>> h = Heap()
182+
>>> h.insert(10)
183+
>>> h.insert(10.1)
184+
>>> h
185+
[10.1, 10]
186+
187+
>>> h = Heap()
188+
>>> h.insert(0.1)
189+
>>> h.insert(0)
190+
>>> h.insert(9)
191+
>>> h.insert(5)
192+
>>> h
193+
[9, 5, 0.1, 0]
194+
"""
126195
self.h.append(value)
127196
idx = (self.heap_size - 1) // 2
128197
self.heap_size += 1

0 commit comments

Comments
 (0)