How to implement Heap in Swift

(If you’re not familiar with Heap, here is the introduction)

When you’re preparing for the coding interview or doing the Leetcode weekly contest, you sometimes need to use Heap instead of sorting to solve the quiz, otherwise, your algorithm may go over the time limitation. (e.g. 295. Find Median from Data Stream)

Unfortunately, in Swift, Apple doesn’t provide Heap for us. That used to be annoying me, but not now. Let’s make a simple Heap on our own, to save you from trouble next time!

1. Declare the fundamental function

Firstly, we need to declare the fundamental function for Heap(the example is min Heap, so the value of parent node must be smaller than the value of child nodes):

Public

1. peek: Return the root node
2. push: Add a new node into the heap
3. pop: Remove the root node from the heap

Private

1. heapifyUp: If the child node is smaller than the parent, which is against the rule, exchange the child and parent. The child will keep shifting up until it fits the rule.
2. heapifyDown: If the parent node is greater than the child, which is against the rule, exchange the parent and child. The parent will keep shifting down until it fits the rule.

Min Heap

2. Find the node

Before implementing the function, we need to use the array to collect the nodes. If we know that the index of the parent node is n, the index of the left child node will be 2n + 1, and the right one will be 2n + 2. On the other hand, if we know that the index of the child node is n, the index of the parent node will be always (n-1) / 2.

3. Implement push and heapifyUp

When we add a new node into the heap, we will append it in the last position of the array, then compare it and its parent. If the new node is smaller than the parent, switch both positions. Keep comparing and switching the new node and its parent until it won’t be smaller than its parent.

Every time you push a new value into the heap, the complexity is O(logN).

4. Implement pop and heapifyDown

When we pop the root node from the heap, the time complexity is O(N) if we remove it from the beginning. We can switch the root and the last node, then call removeLast() or popLast() which can reduce the time complexity of removing root from the array.

But now the last element is on the top, it’s greater than its child nodes, it is against the min-heap rule, so we need to shift it down. Comparing the parent and the child nodes, if the parent is greater than the child nodes, pick the smaller child node then switch it with the parent. Keep doing this until the parent nodes are always smaller than their child nodes.

Here is the whole sample code.

Of course, you can improve it with the generic types, so you don’t need to declare MinHeap and MaxHeap twice.

That’s it. If you face the situation wherein you need to use heap next time, hope you won’t be frustrated anymore.