Minimum Spanning Trees - Prim's Algorithm

MSTs are widely used to calculate optimal paths in a lot of different fields. From a post-office calculating the optimal path for a postman that needs to cover a certain area, to large-scale telecommunication companies finding the cheapest routes for laying down communication cables.

Thanks to constantly increasing computational power, you can manipulate graphs easier than ever. The study of data structures and algorithms has given us better insight into the use of graphs, thanks to algorithms created by great minds over the years.

One of these algorithms is Prim’s Algorithm. It was initially designed by Vojtech Jarnik in 1930, after that Robert C. Prim redesigned it in 1957. This algorithm is used to find the Minimum Spanning Tree in a weighted, undirected graph.

In this lesson, you'll learn how to implement Prim's algorithm in Python - we'll cover the steps used in the algorithm and implement them through a functional example. Before that, we'll dive into the definitions of concepts you need to understand before you try to implement the algorithm itself - trees, graphs, adjacency matrices, and minimum spanning trees.

What Is a (Minimum Spanning) Tree?

A tree is a type of graph, though not all graphs are trees. The main feature of a tree is that every pair of nodes is connected by only one path, so unlike other types of graphs there cannot be cycles in the graph - they're acyclic.

Additionally, trees are undirected - there are no strict directions between nodes that you have to follow. The following illustration will help you get an intuitive understanding of what a tree actually is:

The left graph is a tree and the right one isn't because it has a cycle. Although this might not look much like a tree that you're used to seeing in the park - a lot of trees actually do resemble them, having a hierarchy from the root node with many branches and leaves. Though, to be fair - most trees in computer science are upside down, with the root at the top.

Finding Minimum Spanning Tree (MST) is a common occurrence in the real world. It usually comes in handy when you would like to find the "cheapest" coverage of a certain area. For example, a pizza restaurant chain might want to know where to open up their restaurants to cover the largest area of delivery, within their guaranteed time, with the minimal number of open restaurants - serving the maximum number of people, with the minimum investment, resulting in the biggest return. The result is a minimum spanning tree of their restaurants.

The same analogy can be transferred to various other domains, including delivery networks, telecommunication networks, electrical grids, water supply networks, and almost any infrastructural network.

You can calculate a minimum spanning tree for every undirected graph that has weights assigned to its edges. Each of these trees fulfills all of the following conditions:

• Is a subgraph (this means the MST contains some or all the relationships from the original graph, no more)
• Is a tree, which implies that it has no cycles
• The MST weight (sum of weights) is the minimum weight possible of the different potential spanning trees in the graph.

A spanning tree (maximum or minimum) connects all of the general graph's nodes, but not necessarily all the edges of the graph - some are avoided to reduce cost, and even if not, using all the edges may turn the path into a cyclic one, which then defeats the purpose of a tree.

What is an Adjacency Matrix?

Another concept you need to understand before diving into the implementation of Prim's algorithm is the data structure used to represent a graph. Relationships in a graph can be represented in many ways but the one that is particularly interesting when implementing Prim's algorithm is an adjacency matrix.

This structure shows which nodes are connected to which, meaning that it illustrates the edge structure of a graph. An adjacency matrix is a square matrix with dimensions n x n, where n is the number of nodes in the graph. Once the matrix structure is defined, its fields define which nodes have paths connected to other nodes.

Here's our example tree, re-represented as an adjacency matrix:

We've labeled each node in the tree with a letter. Also, we've labeled each row and column of the corresponding adjacency matrix by one node. Therefore, each field of the adjacency matrix corresponds to two nodes. For example, if a tree has an edge between nodes A and B, the field [A][B] of an adjacency matrix is marked with 1. On the other hand, if there is no edge between those nodes, the field [A][B] is marked with 0.

If you take another look at the example tree, you will notice that it is not directed, meaning that you can traverse each edge from both directions. The adjacency matrix illustrates that by marking both [A][B] and [B][A] when there is an edge between nodes A and B.

That's why an adjacency matrix of a tree is a symmetric matrix. You can notice a distinct diagonal (marked in the illustration above) dividing the matrix into two mirrored sections - called upper and lower triangles.

With this summary, we can get into Prim's Algorithm.

Prim's Algorithm Intuition

In 1957 Robert C. Prim designed (or rather, redesigned) a sequence of steps to find a graph's Minimum Spanning Tree using path weights. The algorithm's steps are these:

1. Select a random node.
2. Choose the path with the minimum weight connected to the chosen node.
3. The path will lead you to a new node, position yourself there.
4. Once you have formed/updated the initial tree, choose the path with the minimum weight that is connected to the whole tree. Keep in mind that you must avoid creating cycles.
5. Repeat steps 3 and 4 until you have covered all the vertices.

Let's take a look at the example graph and find its minimum spanning tree using Prim's algorithm. This graph is the same as the one used in lessons regarding two other algorithms for finding MST in a graph - Kruskal's and Borůvka's:

Firstly, you have to choose a random node to start an algorithm. In this example, let's say you opted to start from node A.

After you choose the starting node, take a look at all edges that are connected to it. In this case, there are two such edges - AB and AC. Then you choose the one with the smaller weight - AB and try to add it to a resulting tree. If it forms a cycle, you skip it and choose an edge with the second minimal weight, etc.

In this case, the edge AB doesn't form any cycles, so you can add it to a resulting tree:

Now, you have a tree consisting of nodes A and B, and the edge connecting those two nodes. Again, you have to take a look at all edges connected to nodes of the tree you formed up until now. The one that has the minimum weight is the edge AC. Since it doesn't create any cycles, you can add it to the tree.

After that, you choose a new edge to add and so on. Let's take a look at this interesting situation that occurs after a couple of iterations:

If you look at all edges connected to the resulting tree, you see that the one with the minimal weight is EH. But if you try to add it to a resulting tree, you will see that it forms a cycle. Therefore, you must skip it and choose the edge with the second minimal weight - DG.

After repeating the steps of the algorithm a few more times, you will get the minimum spanning tree of a resulting graph:

This process can really be appreciated through an animation:

How to Implement Prim's Algorithm in Python

In this section, we'll label nodes of the example graph with numbers instead of letters. That will help us implement the algorithm more easily:

The first thing you need to implement for Prim's algorithm is a Graph class. It is a Python class you will use to represent a graph and all related methods helping you manipulate graphs. In this article, we'll use its simple implementation and alter it a bit to make it more compatible with Prim's algorithm:

class Graph:
    def __init__(self, num_of_nodes):
        self.m_num_of_nodes = num_of_nodes
        # Initialize the adjacency matrix with zeros
        self.m_graph = [[0 for column in range(num_of_nodes)] 
                    for row in range(num_of_nodes)]

    def add_edge(self, node1, node2, weight):
        self.m_graph[node1][node2] = weight
        self.m_graph[node2][node1] = weight

As you can see, this class has a constructor and one method for adding edges to a graph. In this example, we've tweaked the generic constructor it to store a number of nodes in a graph and an adjacency matrix that represents a graph itself. When called, the constructor will initialize the adjacency matrix with all elements set to zero.

The add_edge() method adds an edge to an adjacency matrix. Since Prim's algorithm works with undirected graphs, add_edge() ensures that the resulting matrix is symmetrical.

Prim's Minimum Spanning Tree Algorithm

Now, you can define the prims_mst() method inside of the Graph class. You will use it to define all steps from Prim's algorithm and thus it will yield the MST as the result.

Since you are going to compare weights and look for the minimum one for the starting node, you should define number that works as a temporary minimum before the first node is assigned to the MST. It helps to have this number ridiculously high, such as infinity, to guarantee that the first node we find is of lesser weight than it and that it'll be selected. That's why we'll define a positive_inf for good measure - though, in our example where the numbers are guaranteed to be below 10 - setting the temporary value at 10 would been technically as valid.

We have to track selected nodes to discern which ones are included in the MST. Once all the nodes are a part of the subgraph, you can stop searching for the MST! To achieve this, we are going to create another comprehension list with Boolean values - selected_nodes. Every column in this new comprehension list represents a node. If the node was chosen as part of the MST the field will be True, and False otherwise.

The result will store the minimum spanning tree that is a result of Prim's algorithm. All of these come together in the main section of the algorithm - the while(False in selected_nodes) loop, in which we loop through all of the nodes that haven't yet been selected.

An MST for this purpose doesn't really have a start or end - though, algorithmically, it'll help to have a start and end node. The start will simply be the first randomly selected node, and the end will be the last node we add to the MST. These variables act as the nodes to be connected and we are going to use them to fill in our MST matrix:

def prims_mst(self):
    # Defining a really big number, that'll always be the highest weight in comparisons
    postitive_inf = float('inf')

    # This is a list showing which nodes are already selected 
    # so we don't pick the same node twice and we can actually know when stop looking
    selected_nodes = [False for node in range(self.m_num_of_nodes)]

    # Matrix of the resulting MST
    result = [[0 for column in range(self.m_num_of_nodes)] 
                for row in range(self.m_num_of_nodes)]
    
    indx = 0
    for i in range(self.m_num_of_nodes):
        print(self.m_graph[i])
    
    print(selected_nodes)

    # While there are nodes that are not included in the MST, keep looking:
    while(False in selected_nodes):
        # We use the big number we created before as the possible minimum weight
        minimum = postitive_inf

        # The starting node
        start = 0

        # The ending node
        end = 0

        for i in range(self.m_num_of_nodes):
            # If the node is part of the MST, look its relationships
            if selected_nodes[i]:
                for j in range(self.m_num_of_nodes):
                    # If the analyzed node have a path to the ending node AND its not included in the MST (to avoid cycles)
                    if (not selected_nodes[j] and self.m_graph[i][j]>0):  
                        # If the weight path analized is less than the minimum of the MST
                        if self.m_graph[i][j] < minimum:
                            # Defines the new minimum weight, the starting vertex and the ending vertex
                            minimum = self.m_graph[i][j]
                            start, end = i, j
        
        # Since we added the ending vertex to the MST, it's already selected:
        selected_nodes[end] = True

        # Filling the MST Adjacency Matrix fields:
        result[start][end] = minimum
        
        if minimum == postitive_inf:
            result[start][end] = 0

        print("(%d.) %d - %d: %d" % (indx, start, end, result[start][end]))
        indx += 1
        
        result[end][start] = result[start][end]

    # Print the resulting MST
    # for node1, node2, weight in result:
    for i in range(len(result)):
        for j in range(0+i, len(result)):
            if result[i][j] != 0:
                print("%d - %d: %d" % (i, j, result[i][j]))

Here, we've moved through the adjacency matrix of the initial graph, using two loops. The first loop is for the X-axis (rows) and the second loop is for the Y-axis (columns). Before entering the second loop, you have to validate that the node given by the first loop is selected, which ensures it's part of the MST graph. We handle this with the if selected_nodes[i]: block in the code above.

When you start building the tree, none of the nodes are initially selected, they are all False, so the first loop would end before we even enter the second loop. For this reason, start and end are initially set to 0, and when we exit from the loop, the Boolean value assigned to the end position will become True. As a result, one field of the result will be filled with the existing minimum, and since the result is symmetrical we can use the same trick on the self.m_graph to fill in another field.

Now that we have a selected node, which is Step 1 of Prim's Algorithm, we get to Step 2, within the second loop! First, you move through each column and examine the relationships between our selected node and other nodes. Our selected node will be compared with the other n nodes according to the following parameters:

• The vertex given by i must have a path that connects it with vertex j (this means the weight in the (i,j) position of the adjacency matrix must be greater than zero).
• The vertex j must be not selected (if it's already selected this can lead to a cycle).

Given these two conditions, you can compare the edge weight of a given relationship with the general minimum of the MST. If the weight is less than the minimum, then it will become the new minimum, and the variables start and end will receive the i and j values. If the weight is more than the minimum, then you keep searching through the remaining columns.

The start and end will populate the MST matrix, creating the tree we are looking for. After that, you repeat described process until you select all nodes from the initial graph.

Testing Prim's Algorithm On the Example Graph

To test the previously described implementation of Prim's algorithm, let's create a Graph instance with 9 nodes:

# Example graph has 9 nodes
example_graph = Graph(9)

Then, let's recreate the example graph we've been using earlier in the illustrations and animation:

example_graph.add_edge(0, 1, 4)
example_graph.add_edge(0, 2, 7)
example_graph.add_edge(1, 2, 11)
example_graph.add_edge(1, 3, 9)
example_graph.add_edge(1, 5, 20)
example_graph.add_edge(2, 5, 1)
example_graph.add_edge(3, 6, 6)
example_graph.add_edge(3, 4, 2)
example_graph.add_edge(4, 6, 10)
example_graph.add_edge(4, 8, 15)
example_graph.add_edge(4, 7, 5)
example_graph.add_edge(4, 5, 1)
example_graph.add_edge(5, 7, 3)
example_graph.add_edge(6, 8, 5)
example_graph.add_edge(7, 8, 12)

And, finally, run the algorithm:

example_graph.prims_mst()

That will output the following:

0 - 1: 4
0 - 2: 7
2 - 5: 1
3 - 4: 2
3 - 6: 6
4 - 5: 1
5 - 7: 3
6 - 8: 5

And that's it! That's our MST for that graph - the same one as in the section on Prim's Algorithm Intuition. Each row of the output represent one edge of the resulting MST in the form of node1 - node2: weight:

Remember that we can start with a random node, we don't necessarily need to start with the first one. If you want to challenge yourself, you can modify the code so that it takes a random number (in the correct range of course) as the starting node and observe the algorithm find the same tree in a different order.

What Is the Complexity of Prim's Algorithm?

If you implement Prim's algorithm using only an adjacency matrix, its time complexity is O(N²) - where N is the number of nodes in a graph. The simple implementation yields high time complexity in this case.

You can improve time complexity by opting in for more complex implementation, consisting of Fibonacci or a Binary Heap alongside the adjacency matrix. In that case, the time complexity could be O(E log N), meaning that Prim's algorithm can run as fast as Kruskal's and Borůvka's!

Conclusion

Prim's algorithm is not only efficient but flexible when it comes to finding the Minimum Spanning Tree of a graph. The Python implementation is also really simple. MSTs are useful structures that can be applied in a wide variety of fields, making Prim's algorithm an incredibly important one.