Writing a Heap from scratch in Swift

Why make this post?

A Heap has wide applications when writing highly-performant code, for example in HeapSort, and as such has native implementations in Java and other languages as a "Priority Queue" type. While swift-collections may implement a Heap in the future, I plan on developing a basic Binary Heap using the Swift 5.5 stdlib. Heaps are useful where we need efficient calculation of the minimum/maximum value in a collection.

What is a heap?

A heap is a directed acyclic graph (DAG) structure. i.e.

A graph where edges travel in one direction only, and there are no cycles.

A heap is a tree with a single root node where the whole tree satisfies the heap property.

Heap property

A heap has the following property:

Priority of any node >= Priority of its children

Typically, heaps can have multiple children, however I will focus here on Binary Heaps where each node has at most 2 children.

NOTE: A Priority Queue is implemented as a heap where the root node has the highest priority.

Max Heap vs Min Heap

A Max Heap stores the maximum value at the root node wherease the Min Heap stores the minimum value at the root node.

Representing a Heap in Swift

We can represent a Heap using an underlying Swift Array.

We don't need any additional data structures in this implementation because we can get the children using an indexing strategy as follows:

  1. Indexing starts at 1, with an empty node at index 0. The root node is at index 1.
  2. The Left child of a node is found at the node number doubled.
  3. The Right child of a node is found at the node number doubled + 1.
  4. The Parent of a node is the floor of the node number divided by 2.
var heap = [nil, 1, 2, 3, 4, 5, 6, 7]

What functions does a Heap implement?

The heap has several fundamental operations that distinguish it from other data structures.

  1. insert(_ elem: Int) - O(logN)
  2. heapifyUp(_ elem: Int) - O(logN)
  3. heapifyDown(_ elem: Int) - O(logN)
  4. extractMax() -> Int - O(logN)

insert inserts an element into the Heap while maintaining the Heap property.

heapifyUp (Swim) increases the priority of an element while maintaining the Heap property.

heapifyDown (Sink) decreases the priority of an element while maintaining the Heap property.

extractMax removes the root node from the heap and restructures the Heap maintaining the Heap property.

NOTE: For a Min Heap, this operation would be called extractMin() -> Int.

We will focus here on insert(_ elem: Int) and extractMax() -> Int.

insert(_ elem: Int)

We are essentially inserting the element into the rightmost position into the heap, and bubbling up (swim) until we get a valid heap property.

class MaxHeap {
  private var heap: [Int?] = [nil]
  
  func insert(_ elem: Int) {
		heap.append(elem)
    var nodeNumber = heap.count - 1
    while nodeNumber > 0 {
      if let parent = heap[nodeNumber / 2], 
      	 let child = heap[nodeNumber],
      	 child > parent {
        heap.swapAt(nodeNumber, nodeNumber / 2)
        nodeNumber = nodeNumber / 2
      } else {
        return
      }
    }
  }
}

extractMax() -> Int

We are essentially swapping the root node element with the rightmost position into the heap, removing it, and bubbling the new root node down (sink) until we get a valid heap property.

class MaxHeap {
  private var heap: [Int?] = [nil]
  
  func insert(_ elem: Int) {
    ...
  }
  
  func extractMax() -> Int? {
    guard heap.count > 1 else { return nil }
    heap.swapAt(heap.endIndex - 1, 1)
    let max = heap.removeLast()
    var vIndex = 1
    while 2 * vIndex < heap.endIndex {
      guard let parent = heap[vIndex] else { return max }
      let left = heap[2 * vIndex] ?? Int.min
      var right = Int.min
      if (2 * vIndex + 1) < heap.endIndex {
        right = heap[2 * vIndex + 1] ?? Int.min
      }
      if parent > left, parent > right {
        return max
      }
      var localMax = 2 * vIndex
      if right > left {
        localMax = 2 * vIndex + 1
      }
      heap.swapAt(vIndex, localMax)
      vIndex = localMax
    }
    return max
  }
}

HeapSort extension

In order to derive the sorted array of a Heap, we can achieve this with O(NlogN) Time Complexity by calling extractMax N times.

func heapSort() -> [Int] {
  var sorted: [Int] =[]
  while let max = extractMax() {
    sorted.append(max)
  }
  sorted.reverse()
  return sorted
}

Note that we are using .reverse() instead of prepending because prepends are O(n) and thus we achieve better time complexity by modifying the final result.

We can simply omit the .reverse() call if we want a reverse-sorted array.

Summary

Heaps are messy to write in Swift, especially if we want to avoid using recursive functions. Recursive functions clean up the code, however they require auxiliary space on the call stack to capture the activation records for each function call. As such, I've opted for an iterative approach above.

Happy coding! :D