Kruskal’s Algorithm in Python (with Pseudocode)

Data structures and algorithms are a cornerstone of computer science. In our journey so far, we’ve looked at basic data structures like stacks, queues, and dequeues, linked lists and binary trees, and algorithms like sorting algorithms, tree algorithms, and Dijsktra’s algorithm. Now, let’s take a look at another important graph algorithm – Kruskal’s.

Kruskal’s algorithm is an algorithm that finds the minimum spanning tree (MST) of a graph. In this post we’ll cover:

• Kruskal’s Algorithm Pseudocode
• Kruskal’s Algorithm in Python
  • Creating a Graph Object
  • Python Union Find Algorithm
  • Implementing Kruskal’s Algorithm in Python
  • Full Python Code for Kruskal’s Graph Algorithm
  • Testing Kruskal’s Algorithm in Python
• Summary of Kruskal’s Algorithm in Python

Kruskal’s Algorithm Pseudocode

Kruskal’s algorithm uses a greedy approach to build a minimum spanning tree. Let’s take a look at the pseudocode:

1. Initialize a graph using the shortest (lowest weight) edge
2. Find the shortest connected edge and add it to the shortest edges so far as long as adding the edge doesn’t create a cycle in the graph
3. Repeat step 2 until all vertices have been included in the final MST

Kruskal’s Algorithm in Python

Let’s go over how to implement Kruskal’s Algorithm in Python. This implementation of Kruskal’s Algorithm is going to be as a function in a Graph object. We’ll create a Graph object that will hold the number of vertices in the graph as well as an adjacency list that represents the graph. 

Each entry in the adjacency list will have three entries, the two vertices and the weight of the edge between them. We also need to create functions to perform the find and union pieces of the union find algorithm.

Creating a Graph Object in Python

The first thing we’ll need to do is create our Graph object. All the functions we define later in the tutorial are a part of this object. We need an __init__ function and a function to add an edge. The __init__ function takes one parameter, the number of vertices. It sets one of the object properties to the number of vertices and creates an empty adjacency list to represent the edges in the graph.

class Graph(object):
   def __init__(self, num_vertices):
       self.V = num_vertices
       self.graph = []
  
   def add_edge(self, u, v, w):
       self.graph.append([u, v, w])

Python Union Find Algorithm

Union Find is an algorithm to find the sets of connected graphs. This algorithm consists of two separate functions, the find and union functions. The find function needs two parameters, and the union function needs four parameters.

The find function takes a list that keeps track of the minimum spanning tree and a vertex number. If the index in the MST list is equivalent to the passed in vertex, then we return the vertex number. Otherwise, we recursively look for the vertex by calling the find function with the same list, but with the new index corresponding to the value of the index of the vertex in the MST list.

The union function’s four parameters are the MST list, another list that marks the number of disjoint sets and the origin of those sets, and two vertices. First, we look for where the vertices are in terms of the MST. Then, we compare the values of the indices of the set tracking list and change the corresponding value in the MST list depending on the value comparison. If the values in the MST list are different, then the index of the smaller value is set to the higher value and vice versa. If they are the same, then we create a new disjoint set.

   # union find
   def find(self, root, i):
       if root[i] == i:
           return i
       return self.find(root, root[i])
  
   def union(self, root, rank, x, y):
       xroot = self.find(root, x)
       yroot = self.find(root, y)
       if rank[xroot] < rank[yroot]:
           root[xroot] = yroot
       elif rank[xroot] > rank[yroot]:
           root[yroot] = xroot
       else:
           root[yroot] = xroot
           rank[xroot] += 1

Implementing Kruskal’s Algorithm in Python

Now we have our helper functions, let’s take a look at how we can implement Kruskal’s algorithm in Python. We already have the two objects we need to perform Kruskal’s algorithm so this function does not need any parameters.

The first thing we’ll need to do in our Kruskal’s function is initialize the results list, which will contain the MST and will be structured the same as the list of the vertices and edge weights. We also need to initialize the iterations and number of edges added to the MST. Then, we’ll sort the adjacency list representing the graph by the weights of the algorithm.

Now, we’ll initialize the lists that will represent the root nodes of the sets we’ve found (root) and the list of the number of connected sets in the graph represented by the root node (rank). We’ll then populate the two lists.

Next, we’ll run a while loop that runs as long as we haven’t added the expected number of edges to the MST. In the while loop, we’ll start by extracting the vertices and edge weight between them corresponding to the graph index of the iteration we’re in. Now, we’ll use the find function to see if either or both vertices are connected to the graph.

If only one vertex is connected to the graph (x is not equal to y), then we increment the number of edges we’ve found in the MST, add the set of vertices and edge weights to the results list, and run the union algorithm to keep track of the MST and number of sets. When the MST is created, we’ll print it out.

   # applying kruskal's
   def kruskals(self):
       # initialize an empty MST list
       result = []
       # initialize i, the iteration and e, the edges added
       i, e  = 0, 0
       # sort the graph based on edge weights
       self.graph = sorted(self.graph, key = lambda item: item[2])
       # initialize root, which keeps track of the MST
       # and the rank, which keeps track of where each node belongs
       root = []
       rank = []
       for node in range(self.V):
           root.append(node)
           rank.append(0)
 
       # while we haven't yet added each edge
       # increment iterator and run the union find algorithm
       while e < self.V - 1:
           u, v, w = self.graph[i]
           i = i + 1
           x = self.find(root, u)
           y = self.find(root, v)
           print(f"x, y: {x}, {y}")
           if x != y:
               e = e + 1
               result.append([u, v, w])
               self.union(root, rank, x, y)
 
       for u, v, w in result:
           print(f'{u} - {v}: {w}')

Full Kruskal’s Graph Algorithm Python Code

Here’s the full code for Kruskal’s Algorithm in Python implemented with a Graph class:

# kruskal's in Python
class Graph(object):
   def __init__(self, num_vertices):
       self.V = num_vertices
       self.graph = []
  
   def add_edge(self, u, v, w):
       self.graph.append([u, v, w])
  
   # union find
   def find(self, root, i):
       if root[i] == i:
           return i
       print(i, root[i])
       return self.find(root, root[i])
  
   def union(self, root, rank, x, y):
       print(f"root: {root}, rank: {rank}")
       xroot = self.find(root, x)
       yroot = self.find(root, y)
       if rank[xroot] < rank[yroot]:
           root[xroot] = yroot
       elif rank[xroot] > rank[yroot]:
           root[yroot] = xroot
       else:
           root[yroot] = xroot
           rank[xroot] += 1
       print(f"root: {root}, rank: {rank}")
  
   # applying kruskal's
   def kruskals(self):
       # initialize an empty MST list
       result = []
       # initialize i, the iteration and e, the edges added
       i, e  = 0, 0
       # sort the graph based on edge weights
       self.graph = sorted(self.graph, key = lambda item: item[2])
       # initialize root, which keeps track of the MST
       # and the rank, which keeps track of where each node belongs
       root = []
       rank = []
       for node in range(self.V):
           root.append(node)
           rank.append(0)
 
       # while we haven't yet added each edge
       # increment iterator and run the union find algorithm
       while e < self.V - 1:
           u, v, w = self.graph[i]
           i = i + 1
           x = self.find(root, u)
           y = self.find(root, v)
           print(f"x, y: {x}, {y}")
           if x != y:
               e = e + 1
               result.append([u, v, w])
               self.union(root, rank, x, y)
 
       for u, v, w in result:
           print(f'{u} - {v}: {w}')

Testing Kruskal’s Algorithm in Python

Now let’s test Kruskal’s Algorithm. We’ll create the following graph:

g = Graph(6)
g.add_edge(0, 1, 4)
g.add_edge(0, 2, 4)
g.add_edge(1, 2, 2)
g.add_edge(1, 0, 4)
g.add_edge(2, 0, 4)
g.add_edge(2, 1, 2)
g.add_edge(2, 3, 3)
g.add_edge(2, 5, 2)
g.add_edge(2, 4, 4)
g.add_edge(3, 2, 3)
g.add_edge(3, 4, 3)
g.add_edge(4, 2, 4)
g.add_edge(4, 3, 3)
g.add_edge(5, 2, 2)
g.add_edge(5, 4, 3)
g.kruskals()

This will print out something like the image below.

The MST should correspond to this graph:

Summary of Kruskal’s Algorithm in Python

Kruskal’s algorithm builds a minimum spanning tree of a graph through adding the minimum weighted edges in a greedy manner. We implemented Kruskal’s algorithm in Python by creating a Graph object and adding functions to it to use the union find algorithm to run Kruskal’s algorithm.

Further Reading

• Nested Lists in Python
• Long Short Term Memory (LSTM) in Keras
• The Best Way to do Named Entity Recognition (NER)
• Web Scraping the Easy Way: Python, Selenium, Beautiful Soup 4
• AI Text Summarizer in Python

I run this site to help you and others like you find cool projects and practice software skills. If this is helpful for you and you enjoy your ad free site, please help fund this site by donating below! If you can’t donate right now, please think of us next time.