/ 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