/ eg_heap.py
eg_heap.py
1 2 3 class MinHeap: 4 """ 5 A min-heap of integers 6 """ 7 8 def __init__(self, array=None): 9 """ 10 :param array: of (unordered) elements to build the heap 11 """ 12 self.array = array or [] 13 # BUILD HEAP 14 if len(self.array) > 1: 15 for idx in reversed(range(len(array))): 16 self.heapify(idx) 17 18 19 def find_min(self) -> int: 20 return self.array[0] 21 22 23 def extract_min(self) -> int: 24 # Move last elt to root and heapify root 25 ret = self.array[0] 26 self.array[0] = self.array.pop() 27 self.heapify(0) 28 return ret 29 30 31 def insert(self, x): 32 # Insert at bottom right and bubble up 33 self.array.append(x) 34 cur_idx = len(self.array) - 1 35 par_idx = MinHeap.parent(cur_idx) 36 37 while cur_idx > 0 and self.array[cur_idx] < self.array[par_idx]: 38 self.array[cur_idx], self.array[par_idx] = self.array[par_idx], self.array[cur_idx] 39 cur_idx = par_idx 40 par_idx = MinHeap.parent(cur_idx) 41 42 43 "======== Helpers ========" 44 45 def heapify(self, idx): 46 # Assumes children of idx are valid heaps, but node idx is not necessarily valid 47 # Repeatedly "trickles down": swaps idx with smallest child until idx is valid node 48 left = MinHeap.left(idx) 49 right = MinHeap.right(idx) 50 while (left < len(self.array) and self.array[left] < self.array[idx])\ 51 or (right < len(self.array) and self.array[right] < self.array[idx]): 52 53 new_par = left if right >= len(self.array) or self.array[left] < self.array[right] else right 54 self.array[idx], self.array[new_par] = self.array[new_par], self.array[idx] 55 56 idx = new_par 57 left = MinHeap.left(idx) 58 right = MinHeap.right(idx) 59 60 61 @staticmethod 62 def left(idx): 63 return (idx + 1) * 2 - 1 64 65 66 @staticmethod 67 def right(idx): 68 return (idx + 1) * 2 69 70 71 @staticmethod 72 def parent(idx): 73 return (idx - 1) // 2 74 75 76 77 "========= Testing Heap =========" 78 79 heap = MinHeap([10, 7, 3, 6, 0, 5, 4, 2]) 80 print(heap.array) 81 82 heap.insert(8) 83 print(heap.extract_min()) # pops 0 84 heap.insert(1) 85 heap.insert(9) 86 heap.insert(0) 87 88 print(heap.array) 89 90