Sorting Algorithms: Implementing Heap Sort Using Swift

Photo by Volkan Olmez on Unsplash

Sorting data is a task that has been (and probably will always be) the main concern for computers. Not because sorting in itself is a very interesting topic, but because other algorithms rely on sorted data to function properly. Today, we are going to look at how we can implement a sorting algorithm commonly known as Heap Sort, that relies on the principles of a data structure commonly known as a Heap.

XCode Playground Files with implementations are available on this link.

Revisiting the Heap

A Heap is a binary tree structure that is complete and partially ordered. What this means in normal people language is that it always inserts new data nodes as far to the left in the highest level of nodes as possible. It also means that there is an order to the elements in some sense, even though it is not sorted.

In this tutorial, we are going to use a Max Heap to implement a sorting algorithm that runs in O(n*log(n)) time, always. We will start by implementing a specific SortingAlgorithms class with a static function that sorts your array, and then moves on to a more protocol oriented solution.

Three important things to understand about the Max Heap are:

1. It keeps the largest element in the root at all times.
2. It maintains the property that the value of a node is always bigger than those of its child nodes (if there are any).
3. A Heap can generally be implemented in the form of an array, where parent nodes and child nodes of a specific index can be calculated using very simple mathematical formulas. This will make it possible for us to implement an in-place sort that is fast and efficient.

Getting Into The Code

First off, we are going to start with extending Swift’s standard library’s implementation of the Int structure, in order to abstract away the math that we need to do to keep track of array indices for parent and child nodes. It should look something like this:

private extension Int {
  var parent: Int {
    return (self - 1) / 2
  }
  
  var leftChild: Int {
    return (self * 2) + 1
  }
  
  var rightChild: Int {
    return (self * 2) + 2
  }
}

This implementation makes it possible to calculate indices by calling a computed property instead of placing the math straight into the code. The main win is readability. Also note that this is a private extension, meaning that it will not affect the Int structure as a whole, but will only be available within the scope of this file.

Now, let’s pick apart the different steps in which this sort is going to happen.

First off, we are going to build a Max Heap from our array. By using the fact that all new elements are inserted into the end of the heap and then swapped into place, we are able to simulate insertion using a simple for loop. We will start off pretending that our array only has two elements, “heapify” those two elements, and for each iteration in our loop we are going to pretend that a new element was just inserted and re-heapify from there. It’s going to look a little bit like this:

The first two iterations of building our heap. This continues until all elements have been heapified.

The array will not look especially sorted when this operation is done. In fact, it will look like we made it worse. This is because of the way we use the array to store tree nodes. Take a look at this before/after image:

Before and after heapifying our array.

It looks like we just shuffled the elements. But in fact, we’ve just instated a property that we’re going to use to finalize the sort.

By taking the first element and swapping it for the last element in the array, we will end up with the biggest element in the last place. If we then pretend that the array is one element smaller than it was before, and re-heapify the smaller array, we will end up with the biggest element on the first index again. This means that we can perform exactly the same operation again, and eventually, we will end up with a completely sorted array. Ingenious!

First iteration of shrinking the heap into a sorted array. Note that the second largest element is now at index 0, which makes it possible to perform the same type of swap and re-heap again.

The code for our .heapSort(_:) method goes like this, including building and shrinking the heap.

class SortingAlgorithms {
  private init() {}
  
  public static func heapSort<DataType: Comparable>(_ array: inout [DataType]) {
    if array.count < 2 { return }
    buildHeap(&array)
    shrinkHeap(&array)
  }
  
  private static func buildHeap<DataType: Comparable>(_ array: inout [DataType]) {
    for index in 1..<array.count {
      var child = index
      var parent = child.parent
      while child > 0 && array[child] > array[parent] {
        swap(child, with: parent, in: &array)
        child = parent
        parent = child.parent
      }
    }
  }
  
  private static func shrinkHeap<DataType: Comparable>(_ array: inout [DataType]) {
    for index in stride(from: array.count - 1, to: 0, by: -1) {
      swap(0, with: index, in: &array)
      var parent = 0
      var leftChild = parent.leftChild
      var rightChild = parent.rightChild
      while parent < index {
        var maxChild = -1
        if leftChild < index {
          maxChild = leftChild
        } else {
          break
        }
        if rightChild < index && array[rightChild] > array[maxChild] {
          maxChild = rightChild
        }
        guard array[maxChild] > array[parent] else { break }
        
        swap(parent, with: maxChild, in: &array)
        parent = maxChild
        leftChild = parent.leftChild
        rightChild = parent.rightChild
      }
    }
  }
  
  private static func swap<DataType: Comparable>(_ firstIndex: Int, with secondIndex: Int, in array: inout [DataType]) {
    let temp = array[firstIndex]
    array[firstIndex] = array[secondIndex]
    array[secondIndex] = temp
  }
}

Perfect! This does exactly what we want it to! But… can we make it cleaner?

The Protocol Oriented Approach

By extending the Array type for element types that are comparable, we will get a few benefits. One is that we need less code, but the main gain is that we will be able to call the method directly on one of our objects, so instead of doing this:

SortingAlgorithms.heapSort(&myArray)

We will be able to do this:

myArray.heapSort()

This is much cleaner, and the compiler will not even show us the option to call .heapSort() on an array whose elements are not conforming to the Comparable protocol.

public extension Array where Element: Comparable {
  public mutating func heapSort() {
    buildHeap()
    shrinkHeap()
  }
  
  private mutating func buildHeap() {
    for index in 1..<self.count {
      var child = index
      var parent = child.parent
      while child > 0 && self[child] > self[parent] {
        swapAt(child, parent)
        child = parent
        parent = child.parent
      }
    }
  }
  
  private mutating func shrinkHeap() {
    for index in stride(from: self.count - 1, to: 0, by: -1) {
      swapAt(0, index)
      var parent = 0
      var leftChild = parent.leftChild
      var rightChild = parent.rightChild
      while parent < index {
        var maxChild = -1
        if leftChild < index {
          maxChild = leftChild
        } else {
          break
        }
        if rightChild < index && self[rightChild] > self[maxChild] {
          maxChild = rightChild
        }
        guard self[maxChild] > self[parent] else { break }
        
        swapAt(parent, maxChild)
        parent = maxChild
        leftChild = parent.leftChild
        rightChild = parent.rightChild
      }
    }
  }
}

That’s it for this time! Feel free to comment if you have questions, and follow to get notifications about future articles.

To learn more about iOS Development, check out my previous articles:

The Process of Designing a Custom Data Structure

Using Swift Property Observers To Clean Up Objects