DEV Community

Sarvesh Kesharwani
Sarvesh Kesharwani

Posted on

Implementing max heap from scratch in Python!

class MaxHeap:
    def __init__(self):
        self.heap = []

    def insert(self, val):
        self.heap.append(val)
        self._percolate_up(len(self.heap)-1)

    def extract_max(self):
        if len(self.heap) > 1:
            max_val = self.heap[0]
            self.heap[0] = self.heap.pop()
            self._percolate_down(0)
        elif len(self.heap) == 1:
            max_val = self.heap.pop()
        else:
            max_val = None
        return max_val

    def _percolate_up(self, idx):
        parent_idx = (idx - 1) // 2
        if parent_idx < 0:
            return
        if self.heap[idx] > self.heap[parent_idx]:
            self.heap[idx], self.heap[parent_idx] = self.heap[parent_idx], self.heap[idx]
            self._percolate_up(parent_idx)

    def _percolate_down(self, idx):
        child_idx1 = idx * 2 + 1
        child_idx2 = idx * 2 + 2
        max_idx = idx
        if child_idx1 < len(self.heap) and self.heap[child_idx1] > self.heap[max_idx]:
            max_idx = child_idx1
        if child_idx2 < len(self.heap) and self.heap[child_idx2] > self.heap[max_idx]:
            max_idx = child_idx2
        if max_idx != idx:
            self.heap[idx], self.heap[max_idx] = self.heap[max_idx], self.heap[idx]
            self._percolate_down(max_idx)

Enter fullscreen mode Exit fullscreen mode

In the code above, we define a MaxHeap class with methods to insert elements into the heap and extract the maximum value from the heap. The _percolate_up and _percolate_down methods are used to maintain the heap property after inserting or extracting elements.

To create a new MaxHeap object and insert values into it, you can do the following:

heap = MaxHeap()
heap.insert(10)
heap.insert(30)
heap.insert(20)
heap.insert(5)
Enter fullscreen mode Exit fullscreen mode

To extract the maximum value from the heap, you can call the extract_max method:

max_val = heap.extract_max()
print(max_val) # Output: 30
Enter fullscreen mode Exit fullscreen mode

Top comments (0)