Swift Algorithm Club: Minimum Spanning Tree with Prim’s Algorithm

The Swift Algorithm Club is an open source project on implementing data structures and algorithms in Swift.

Every month, Kelvin Lau, Ross O’Brien and I feature a cool data structure or algorithm from the club in a tutorial on this site. If you want to learn more about algorithms and data structures, follow along with us!

In this tutorial you will learn to implement a Swift minimum spanning tree, specifically using prim’s algorithm.

Prim’s algorithm was first implemented for the Swift Algorithm Club by Xiang Xin (thank you!), and has been presented here for tutorial format.

Note: If you have been follow along our SAC article series, last time you saw how to create a heap and priority queue. You will now use a priority queue to implement Prim’s algorithm.

Introduction

Prim’s algorithm was first discovered by a mathematician named Vojtěch Jarník, and later again by Robert Prim. This algorithm is also known as a greedy algorithm.

A greedy algorithm is used to find optimal solutions. It constructs a solution step by step, where at every stage it picks the most optimal path.

Prim’s algorithm constructs a minimum spanning tree. A minimum spanning tree is formed by a subset of connected undirected weighted edges, that connect all vertices together without forming a cycle. Also it has the minimum possible total edge weight. For example you might want to find the cheapest way to layout your water pipes effectively, to cut cost.

In Prim’s algorithm you create a minimum spanning tree by picking edges one at a time. It’s a greedy algorithm because every time you pick an edge, you always pick the smallest weighted edge that connects a pair of vertices.

For other applications of greedy algorithm check out this link by UT Dallas.

Getting Started

Before you start implementing prim’s algorithm, there are 6 steps to consider when performing it:

Let’s go through a simple example. Imagine the network graph below. It could represent any type of network! Let’s open up our imagination here!

Imagine a network of airports represented by the vertices. The weighted edges represents the cost and route for an airplane to fly from one airport to the next. RW Airways may want to adopt prim’s algorithm to get a minimum spanning tree giving them the most cost effective routes to fly between airports!

By the end of this tutorial you will save lots of dollar bills!

Working through an example

1. First, start by picking any vertex. Let’s say vertex 2.
2. Choose the shortest edge from this vertex. This vertex has edges with weights { 6, 5, 3 }.
3. You pick edge 3 since it’s the smallest, which routes to vertex 5.

1. You have now visited vertices { 2, 5 }.
2. Choose the shortest edge from these vertices. The next nearest edges are { 6, 5, 6, 6 }. You pick edge 5 since it’s the smallest, which routes to vertex 3.
3. Notice that edge 6 between vertex 5 and vertex 3 can be removed since these vertices have already been visited.

1. You now have visited vertices { 2, 3, 5 } .
2. Choose the shortest edge from these vertices. The next nearest edges are { 6, 1, 5, 4, 6 } . You pick edge 1 since it’s the smallest, which routes to vertex 1.
3. Notice that edge 6 between vertex 2 and vertex 1 can be removed since these vertices have already been visited.

1. You now have visited vertices { 2, 3, 5, 1 }.
2. Choose the shortest edge from these vertices. The next nearest edges are { 5, 5, 4, 6 } . You pick edge 4 since it’s the smallest, which routes to vertex 6.
3. Notice that edge 6 between vertex 5 and vertex 6 can be removed since these vertices have already been visited.

1. Finally you now have visited vertices {2, 5, 3, 1, 6 }.
2. Choose the shortest edge from these vertices. The next nearest edges are { 5, 5, 2 } You pick edge 2 since it’s the smallest, which routes to vertex 4 .
3. Notice that the rest of the edges { 5, 5 } connected to vertex 4 from vertex 1 and vertex 3 can be removed since these vertices have been visited.

RW Airways can now redirect their airplanes to fly these routes to cut cost! This means more profit! Are you starting to see the pattern here? Let’s move on to the implementation!

Implementation

Download the starter playground. Open the Project Navigator to see the list of files, and familiarize yourself with the Sources folder:

1. You will be using an adjacency list to create your graph network.
2. You will be using a Priority Queue (implemented using a Heap). This data structure will store edges adjacent to vertices you visit. The root of this heap will always have the smallest weighted edge.

Again, if you want more information on how these classes work, check out our Swift heap and priority queue tutorial.

Setup

Navigate to the root Playground and add the following code:

class Prim<T: Hashable> { // 1
  typealias Graph = AdjacencyList<T> // 2
  var priorityQueue = PriorityQueue<(vertex: Vertex<T>, weight: Double, parent: Vertex<T>?)>(
    sort: { $0.weight < $1.weight }) // 3
}

1. You have declared a class named Prim, you want this algorithm to be generic, so it can find the minimum spanning tree of any type, given that it's Hashable.
2. This class will use the AdjacencyList to hold your graphs.
3. You define a minimum priorityQueue, where it holds the current vertex, the weight of the edge between the vertices, and the parent vertex.

Next add the following after priorityQueue:

func produceMinimumSpanningTree(graph: Graph) -> (cost: Double, mst: Graph) { // 1
  var cost = 0.0 // 2
  let mst = Graph() // 3
  var visited = Set<Vertex<T>>() // 4

  return (cost: cost, mst: mst) // 5
}

1. Define a function produceMinimumSpanningTree(_:), that takes a Graph network, and returns the minimum total cost of all edges, and the minimum spanning tree.
2. Define the cost variable, used to aggregate all weights.
3. Create a graph to construct your minimum spanning tree.
4. Create a Set to store all vertices visited.
5. Once prim's algorithm is complete, return the results.