Max heap implementation in Python using heapq

This post will discuss how to implement max heap in Python based on the heapq module.

1. Max Heap of primitives

The heapq module in Python provides the min-heap implementation of the priority queue algorithm. We can easily implement max heap data structure using it.

The following program provides a simple implementation of max heap for integers using heapq operations. It can be easily extended to support any other general-purpose functions based on heaps.

import heapq

class MaxHeap:

# Initialize the max heap

def __init__(self, data=None):

if data is None:

self.data = []

else:

self.data = [-i for i in data]

heapq.heapify(self.data)

# Push item onto max heap, maintaining the heap invariant

def push(self, item):

heapq.heappush(self.data, -item)

# Pop the largest item off the max heap, maintaining the heap invariant

def pop(self):

return -heapq.heappop(self.data)

# Pop and return the current largest value, and add the new item

def replace(self, item):

return heapq.heapreplace(self.data, -item)

# Return the current largest value in the max heap

def top(self):

return -self.data[0]

if __name__ == '__main__':

input = [7, 4, 6, 3, 9, 1]

# build a max heap from all elements in the list

pq = MaxHeap(input)

# pop from max heap

print(pq.pop()) # 9

print(pq.pop()) # 7

print(pq.pop()) # 6

pq.push(10)

pq.push(9)

print(pq.pop()) # 10

print(pq.pop()) # 9

Download Run Code

Output:

9
7
6
10
9

2. Max Heap of Objects

For objects, we can directly use the heapq module in Python to get a max heap. The idea is to override the less-than operator __lt__ to make our class work with max heap. The following program demonstrates this:

import heapq

from heapq import heappop

class Pair:

def __init__(self, x, y):

self.x = x

self.y = y

# Override the less-than operator __lt__ to make Pair class work with max heap

def __lt__(self, other):

return self.x > other.x

def __repr__(self):

return f'({self.x}, {self.y})'

if __name__ == '__main__':

# build a max heap of Pair

pq = [Pair(7, 0), Pair(3, 3), Pair(9, 4), Pair(4, 1), Pair(6, 2), Pair(1, 5)]

heapq.heapify(pq)

# run until max heap is empty

while pq:

# print the next maximum element from the heap

print(heappop(pq))

Download Run Code

Output:

(9, 4)
(7, 0)
(6, 2)
(4, 1)
(3, 3)
(1, 5)

That’s all about max heap implementation in Python using the heapq module.