Prim's Algorithm in Python: A Guide to Efficient Graph Management

Imagine you are managing a project for a major telecommunications company, such as Airtel or Jio, and your task is to lay cable lines underground efficiently and cost-effectively. This scenario is a perfect application of Prim’s algorithm, a fundamental tool in network design and graph management.

In this article, we will understand Prim’s algorithm, how it is solved, and how to implement it in Python.

Prim’s algorithm finds the minimum spanning tree (MST) for a weighted undirected graph. It starts with a single vertex and expands the MST by adding the nearest vertex not yet included. It’s ideal for dense graphs where you want to ensure minimal connection costs.

What is Prim’s algorithm?

Prim’s algorithm, an extension of the greedy approach, helps in finding the minimum spanning tree (MST). Essentially we start with a vertex and then find the shortest edge to the next vertex. This process is repeated until we get our MST. We keep repeating the same process with other vertices and finally, we receive a bunch of MSTs. We choose the most suitable and optimized Minimum Spanning Tree from all the options.

Let us try to implement all this in Python.

Implementing Prim’s Algorithm in Python

Given below we have written the code for Prim’s algorithm. It follows the simple greedy algorithm, and we plot the graph after repetitive processes of finding the least distance between vertices.

import matplotlib.pyplot as plt
import numpy as np

class Graph:
    def __init__(self, vertices):
        self.V = vertices  # No. of vertices
        self.graph = [[] for i in range(self.V)]  # default graph will be filled with an empty list

    def addEdge(self, u, v, w):
        self.graph[u].append((v, w))
        self.graph[v].append((u, w))  # for undirected graph

    def printMST(self, parent):
        print("Edge \tWeight")
        for i in range(1, self.V):
            print(parent[i], "-", i, "\t", self.graph[i][0][1])

    def minKey(self, key, mstSet):
        min = np.inf  # Initialize min value
        min_index = None

        for v in range(self.V):
            if key[v] < min and mstSet[v] == False:
                min = key[v]
                min_index = v

        return min_index

We introduce the required libraries like matplotlib and numpy for our operations. We then create a class called Graph that gives us the structure of a graph where any vertex can connect to any other vertex. Furthermore, we find a minimum spanning tree of the graph using our Prim’s algorithm. The minimum spanning tree is used to connect all the points with the least weights of edges.

    def primMST(self):
        parent = [-1] * self.V  # Array to store constructed MST
        key = [np.inf] * self.V  # Key values used to pick minimum edge in cut
        mstSet = [False] * self.V  # Represents set of vertices included in MST

        key[0] = 0  # Always include first 1st vertex in MST.
        parent[0] = -1  # First node is always root of MST

        for count in range(self.V):
            u = self.minKey(key, mstSet)

            # Add the picked vertex to the MST Set
            mstSet[u] = True

            # Update key and parent index of the adjacent vertices of the picked vertex.
            # Consider only those vertices which are not yet included in MST
            for v, weight in self.graph[u]:
                if mstSet[v] == False and weight < key[v]:
                    parent[v] = u
                    key[v] = weight

        self.printMST(parent)

The primMST function initializes with a set of arrays to track the minimum edges and whether a vertex is included in the MST. It starts with the first vertex and explores its adjacent vertices, updating the minimum edge weight and parent node if a cheaper connection is found. This ensures that each vertex added to the MST is connected through the least expensive edge. The loop continues until all vertices are included in the MST, ensuring an optimal structure.”

# Sample graph
g = Graph(5)
g.addEdge(0, 2, 1)
g.addEdge(0, 1, 2)
g.addEdge(1, 2, 3)
g.addEdge(1, 3, 4)
g.addEdge(2, 3, 5)
g.addEdge(3, 4, 6)

# Prim's Minimum Spanning Tree (MST)
g.primMST()

# Visualize the MST
edges = []
weights = []
for i in range(g.V):
    for neighbor, weight in g.graph[i]:
        if i < neighbor:
            edges.append((i, neighbor))
            weights.append(weight)

pos = np.random.rand(g.V, 2)  # Random node positions
plt.plot(pos[:, 0], pos[:, 1], 'o', markersize=15, color='black')  # Plot nodes
for edge, weight in zip(edges, weights):
    plt.plot(
        [pos[edge[0]][0], pos[edge[1]][0]],
        [pos[edge[0]][1], pos[edge[1]][1]],
        '-',
        color='blue',
        linewidth=2,
    )  # Plot edges
plt.xlabel("X")
plt.ylabel("Y")
plt.title("Prim's Minimum Spanning Tree")
plt.show()

The visual output created by Matplotlib shows the minimum spanning tree generated by Prim’s algorithm. Each node represents a city, and each line between nodes represents a cable line, with weights indicating the cost. The layout is generated randomly for visualization purposes, helping to illustrate how Prim’s algorithm connects cities in the most cost-effective manner.

Prims Algorithm Output — ***Visualizing the Output of Prim’s Algorithm***

From the above code, we can observe a network between the different nodes. As discussed earlier, this node can be replaced by the name of cities for laying down cable lines.

This algorithm typically works on a queue-type basis and finally, a cyclical-type graph will be created as observed in the above output.

Optimizing Infrastructure with Prim’s Algorithm

Now you understand how Prim’s algorithm can be pivotal in reducing infrastructure costs through strategic planning and efficient resource allocation. By applying this algorithm, companies can significantly optimize their network designs. What other sectors could benefit from such graph-based optimization strategies?