Learning Heaps with Python

Photo by Markus Spiske on Unsplash

A Heap is a special Tree-based data structure in which the tree is a complete Binary tree. The Heap data structure is very efficient to find the kth largest or smallest element. Heaps allow you to pull the smallest or the biggest element and quickly know the next smallest or biggest element. That is why it’s called a priority queue.

Each insert and delete operation is O(log N) at the very worst. And here is the best part, you can get either the smallest element or biggest element in Constant time i.e., O(1)

How do we implement it?

There are two ways to implement Heaps,

  1. Using a list of nodes where every node contains two child nodes
  2. Using an array or list. (Reduces the overhead of node)

We typically come across two types of Heaps:

  1. Max-Heap: In a Max-Heap the key present at the root must be greatest among the keys present at all of its children. The same property must be recursively true for all sub-trees in that Binary Tree. This is used to find the kth largest element.
  2. Min-Heap: In a Min-Heap the key present at the root must be smallest among the keys present at all of its children. The same property must be recursively true for all sub-trees in that Binary Tree. This is used to find the kth smallest element

Implementing a Min-Heap using a Python List !!!

class MinHeap():

def __init__(self, size):
self._elements = [0] * size
self._size = 0

The class name is MinHeap and the constructor takes size as a parameter to initialize the size of our Heap. _elements is a list, with all its values made 0. As the heap is empty now, its _size is 0.

def getLeftChildIndex(self, elementIndex):
return 2 * elementIndex + 1

def getRightChildIndex(self, elementIndex):
return 2 * elementIndex + 2

def getParentIndex(self, elementIndex):
return (elementIndex - 1) / 2

The above methods give the Index of its left child, right child, and parent by passing an element Index. The index of a Heap starts from 0 and go on like this:

Indexing of nodes in a Tree (Heap). The node at the top is called “root” of Heap.
def hasLeftChild(self, elementIndex):
return self.getLeftChildIndex(elementIndex) < self._size

def hasRightChild(self, elementIndex):
return self.getRightChildIndex(elementIndex) < self._size

def isRoot(self, elementIndex):
return elementIndex == 0

The above methods would return True or False if Index hasLeftChild, hasRightChild, and isRoot on passing an elementIndex.

def getLeftChild(self, elementIndex):
return self._elements[self.getLeftChildIndex(elementIndex)]

def getRightChild(self, elementIndex):
return self._elements[self.getRightChildIndex(elementIndex)]

def getParent(self, elementIndex):
return self._elements[self.getParentIndex(elementIndex)]

These methods would return the value stored in its leftChild, rightChild, and parent on passing an elementIndex.

def swap(self, firstIndex, secondIndex):
temp = self._elements[firstIndex]
self._elements[firstIndex] = self._elements[secondIndex]
self._elements[secondIndex] = temp

def isEmpty(self):
return self._size == 0

def peek(self):
if self._size == 0:
raise IndexError
return self._elements[0]

Method swap swap’s value stored in firstIndex and SecondIndex. isEmpty returns True if the Heap is empty. peek returns the root/top element of Heap. If the Heap size is 0, then peek raises an IndexError.

def pop(self):
if self._size == 0:
raise IndexError

result = self._elements[0]
self._elements[0] = self._elements[self._size - 1]
self._size -= 1

self.heapifyDown()
return result

pop returns & removes the root/top element from the Heap. It then initializes the root element (i.e. index = 0) to the last element of the Heap. Then we call heapifyDown which rearranges the Min-Heap in order (which brings the smallest element to the top). As an element is removed from the Heap, we decrease its size by 1.

def add(self, element):
if self._size == len(self._elements):
raise IndexError

self._elements[self._size] = element
self._size += 1

self.heapifyUp()

add add’s an element to the Heap. The new element is stored at the end of the list and then we call heapifyUp which ensures that element falls in its correct spot in the Heap. As an element is added to the Heap, we increase its size by 1.

def heapifyDown(self):
index = 0
while(self.hasLeftChild(index)):
smallerIndex = self.getLeftChildIndex(index)
if(self.hasRightChild(index) and self.getRightChild(index) < self.getLeftChild(index)):
smallerIndex = self.getRightChildIndex(index)

if(self._elements[smallerIndex] > self._elements[index]):
break

self.swap(smallerIndex, index)
index = smallerIndex

heapifyDown starts from index 0. smallerIndex stores rightChildIndex if the right child element is lesser than left Child otherwise it stores leftChildIndex. It then swaps smallerIndex element with index element if smallerIndex element value is greater than index element otherwise it breaks the loop. Later, the index is updated to smallerIndex. This continues till child elements are greater than index element. Here is how heapifyDown would look:

heapifyDown
def heapifyUp(self):
index = self._size - 1
while(not self.isRoot(index)
and self._elements[index] < self.getParent(index)):
parentIndex = self.getParentIndex(index)
self.swap(index, parentIndex)
index = parentIndex
def displayHeap(self):
print(str(self._elements))

displayHeap prints the elements of the Heap. heapifyUp starts with index as the last element of Heap. It then swaps with the parent element if index element is less than parent element. index is then updated to parentIndex. This continues until the parent element is no longer greater the index element. Here is how heapifyUp would look:

heapifyUp

If you wanna play around with it (like add, pop, peek and display heap) here is the full code:

"""
Implementation of a Min-Heap.
"""

class MinHeap():

def __init__(self, size):
self._elements = [0] * size
self._size = 0

def getLeftChildIndex(self, elementIndex):
return 2 * elementIndex + 1

def getRightChildIndex(self, elementIndex):
return 2 * elementIndex + 2

def getParentIndex(self, elementIndex):
return (elementIndex - 1) / 2

def hasLeftChild(self, elementIndex):
return self.getLeftChildIndex(elementIndex) < self._size

def hasRightChild(self, elementIndex):
return self.getRightChildIndex(elementIndex) < self._size

def isRoot(self, elementIndex):
return elementIndex == 0

def getLeftChild(self, elementIndex):
return self._elements[self.getLeftChildIndex(elementIndex)]

def getRightChild(self, elementIndex):
return self._elements[self.getRightChildIndex(elementIndex)]

def getParent(self, elementIndex):
return self._elements[self.getParentIndex(elementIndex)]

def swap(self, firstIndex, secondIndex):
temp = self._elements[firstIndex]
self._elements[firstIndex] = self._elements[secondIndex]
self._elements[secondIndex] = temp

def isEmpty(self):
return self._size == 0

def peek(self):
if self._size == 0:
raise IndexError
return self._elements[0]

def pop(self):
if self._size == 0:
raise IndexError

result = self._elements[0]
self._elements[0] = self._elements[self._size - 1]
self._size -= 1

self.heapifyDown()
return result

def add(self, element):
if self._size == len(self._elements):
raise IndexError

self._elements[self._size] = element
self._size += 1

self.heapifyUp()

def heapifyDown(self):
index = 0
while(self.hasLeftChild(index)):
smallerIndex = self.getLeftChildIndex(index)
if(self.hasRightChild(index) and self.getRightChild(index) < self.getLeftChild(index)):
smallerIndex = self.getRightChildIndex(index)

if(self._elements[smallerIndex] > self._elements[index]):
break

self.swap(smallerIndex, index)
index = smallerIndex

def heapifyUp(self):
index = self._size - 1
while(not self.isRoot(index) and self._elements[index] < self.getParent(index)):
parentIndex = self.getParentIndex(index)
self.swap(index, parentIndex)
index = parentIndex

def displayHeap(self):
print(str(self._elements))

def main():
size = int(raw_input("size of heap?\n"))
heap = MinHeap(size)

while True:
choice = int(raw_input("1. add element\t2. pop element\t3. peek element\t4. display heap\t5. exit\n"))
if choice == 1:
element = int(raw_input("enter element?\n"))
heap.add(element)
elif choice == 2:
print('Removed {}'.format(heap.pop()))
elif choice == 3:
print('Top element is {}'.format(heap.peek()))
elif choice == 4:
heap.displayHeap()
elif choice == 5:
break
else
:
continue


if
__name__ == '__main__':
main()

Max-Heap is almost same with some minor tweaks in heapifyUp and heapifyDown methods.

That’s it for now, feel free to share your comments! :)

Note: There is an inbuilt module called ‘heapq’ in Python.

--

--

--

programmer

Love podcasts or audiobooks? Learn on the go with our new app.

Recommended from Medium

How to Speed Up UI Development with KnowCode for Flutter

What is AmRDS?

Making a seamless Perlin noise in C#

Insertion Sort A Linked List

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store
Ashwanth Kumar

Ashwanth Kumar

programmer

More from Medium

Data Structures in Python

Python: Methods

python

Welcome to the “First Program” in Python? #Python Series-4

A Beginner’s Guide to Basic Python