Understand Heaps by solving a popular interview problem in python

Ashutosh Kumar

Aug 24, 2020·2 min read

In Coding Interviews we do not create a heap data structure from scratch, instead, we use a library function to implement heap. In Python heap is implemented using the heapq library. At the end of this article, using this heapq library we will solve a popular interview problem “Sort an array based on its frequency”.

Heaps typically use an array-based data structure, which solves various problems in optimal time. Few Problems where Heap Data Structure is used are:

1. Finding Minimum or Maximum element in an array
2. Finding Median in a stream of numbers
3. Find Kth Largest or Kth Smallest element in an array

Heaps are of two types: Max-Heap and Min-Heap. In Java, both min-heap and max-heap is implemented using a Priority Queue, But in python, the implementation is a bit different as Python’s heapq module does not have built-in support for the max heap but there is a trick to implement max heap.

Min Heap in Python

import heapqarr = [1,5,10, 9, 3]heapq.heapify(arr)// To push an element to heap use heappush(heap_array,element)
heapq.heappush(arr,2)// Display the first 3 minimum elements of the array
for i in range(3):
    print(heapq.heappop(arr))// Output : 1,2,3

Max Heap in Python

Implementing max heap is a little tricky because heapq doesn’t have built-in support that sorts the elements based on maximum element first. By default heapq implements min-heap.

The trick here is to convert the array into negative elements and push it to the heap as shown in the code below

import heapqli = [1,5,10, 9, 3,2]arr=[]
k=3for i in range(len(li):
    arr.append(-li[i])
    
heapq.heapify(arr)// Display the first three aximum elements in the array
for i in range(3):
    print(-(heapq.heappop(arr)))// Output: 10,9,5

Sort an array based on its frequency

Problem Statement: Given an array with repeated elements, Print distinct elements sorted based on its frequency.

Example:

Input: [15,8,6,7,6,10,9,7,8,7]
Output: 7 8 6 15 10 9

Solution:

import heapqarr = [15,8,6,7,6,10,9,7,8,7]heap=[]
heapq.heapify(heap)
d={}// Create a dictionar to store the frequency of elements
for i in arr:
    if i in d:
        d[i]+=1
    else:
        d[i]=1// Push a tuple of frequency and array element
// The heap is sorted based on the first element of tuplefor i,j in d.items():
    heapq.heappush(heap,(j,i))// nlargest function returns the first n largest numbers of heap
result=heapq.nlargest(len(heap),heap)for i,j in result:
    print(j,end=" ")

Heap is an important data structure for problem-solving in coding interviews and competitive programming. Thank you for reading, I hope this article was useful for you.

References:

1. Geeks For Geeks: https://www.geeksforgeeks.org/min-heap-in-python
2. Wikipedia: https://en.wikipedia.org/wiki/Heap_(data_structure)
3. Python Documentation: https://docs.python.org/3/library/heapq.html