Implementing a Heap Data Structure in Python - Data Structures and Algorithms (DSA) Guide

Introduction

A heap is a special tree-based data structure that satisfies the heap property. It is mainly used to implement priority queues.

A heap can be visualized as a binary tree filled on almost all levels, except that the last level might not be completely filled and it’s filled from left to right.

In a Max Heap, for any given node I, the value of I is greater than or equal to the values of its children. In a Min Heap, the value of I is less than or equal to the values of its children.

Structure of a Heap

Here’s a text-based representation of a Max Heap:

And a Min Heap:

Python Implementation

Python provides an in-built module heapq for implementing heaps quickly. But let’s implement a Min Heap from scratch for a better understanding of the underlying mechanics.

Heap Class

We will use a list to store heap elements and perform operations to maintain the heap property.

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

    def parent(self, i):
        return (i - 1) // 2

    def left_child(self, i):
        return 2 * i + 1

    def right_child(self, i):
        return 2 * i + 2

    def has_left_child(self, i):
        return self.left_child(i) < len(self.heap)

    def has_right_child(self, i):
        return self.right_child(i) < len(self.heap)

    def insert(self, key):
        self.heap.append(key)
        self.heapify_up(len(self.heap) - 1)

    def heapify_up(self, i):
        while i > 0 and self.heap[i] < self.heap[self.parent(i)]:
            self.heap[i], self.heap[self.parent(i)] = self.heap[self.parent(i)], self.heap[i]
            i = self.parent(i)

    def extract_min(self):
        if len(self.heap) == 0:
            return None
        min_val = self.heap[0]
        self.heap[0] = self.heap[-1]
        self.heap.pop()
        self.heapify_down(0)
        return min_val

    def heapify_down(self, i):
        min_index = i
        l = self.left_child(i)
        r = self.right_child(i)

        if self.has_left_child(i) and self.heap[l] < self.heap[min_index]:
            min_index = l
        if self.has_right_child(i) and self.heap[r] < self.heap[min_index]:
            min_index = r

        if i != min_index:
            self.heap[i], self.heap[min_index] = self.heap[min_index], self.heap[i]
            self.heapify_down(min_index)

Method Details:

insert: Inserts a new element into the heap while maintaining the heap property.
extract_min: Removes and returns the smallest element from the heap.
heapify_up: Maintains heap property after insertion.
heapify_down: Maintains heap property after deletion.

Using the MinHeap Class

Here’s an example using the MinHeap class:

# Create a new MinHeap
min_heap = MinHeap()

# Insert elements
min_heap.insert(3)
min_heap.insert(1)
min_heap.insert(4)
min_heap.insert(1)
min_heap.insert(5)
min_heap.insert(9)

# Extract the minimum elements
print(min_heap.extract_min())  # Output: 1
print(min_heap.extract_min())  # Output: 1

Using Python’s In-Built Heap

Python’s standard library provides a heapq module for creating binary heaps with ease. The heapq module transforms a regular list into a heap in-place:

import heapq

# Create a heap
my_heap = []
heapq.heapify(my_heap)

# Insert elements
heapq.heappush(my_heap, 3)
heapq.heappush(my_heap, 1)
heapq.heappush(my_heap, 4)

# Extract the minimum
print(heapq.heappop(my_heap))  # Output: 1

If you find this guide helpful, please share with others, and spread the knowledge!