Prim's Algorithm (Simple Implementation Using Adjacency Matrix)

Introduction

Prim's algorithm is a greedy algorithm that finds a minimum spanning tree for a weighted undirected graph. A minimum spanning tree is a subset of the edges that connects all vertices in the graph with the minimum possible total edge weight. Prim's algorithm starts with an arbitrary starting vertex and adds the edge with the smallest weight to the growing tree at each step, until all vertices are included in the tree.

In this article, we will provide a simple implementation of Prim's algorithm using the adjacency matrix representation of a graph.

Algorithm

1. Create a set MST to hold the edges of the minimum spanning tree and a set visited to hold the visited vertices.
2. Randomly select a starting vertex v0 and add it to the visited set.
3. Repeat until all vertices are visited:
   1. For each visited vertex v, find the adjacent vertices that have not been visited and add the edge with the smallest weight to a priority queue pq.
   2. While pq is not empty, remove the edge with the smallest weight and check if its vertices have not been visited yet. If yes, add the edge to the MST set and add the unvisited vertex to the visited set.
4. Return MST.

Implementation

import heapq

def prim(graph):
    n = len(graph)
    MST = set()
    visited = set()
    v0 = 0
    visited.add(v0)
    pq = [(cost, v0, v) for v, cost in enumerate(graph[v0]) if cost > 0]
    heapq.heapify(pq)
    while len(visited) < n:
        cost, u, v = heapq.heappop(pq)
        if v not in visited:
            visited.add(v)
            MST.add((u, v, cost))
            for w, cost in enumerate(graph[v]):
                if cost > 0 and w not in visited:
                    heapq.heappush(pq, (cost, v, w))
    return MST

The prim() function takes a weighted adjacency matrix graph as input and returns the set MST containing the edges of the minimum spanning tree.

The function first initializes the starting vertex v0 to 0 and adds it to the visited set. It then creates a priority queue pq containing the edges adjacent to v0 that have not been visited yet.

The function then enters a loop that continues until all vertices are visited. At each iteration, the function removes the edge with the smallest weight from pq and checks if its second vertex has not been visited yet. If yes, the function adds the edge to the MST set and adds the unvisited vertex to the visited set. It then adds the adjacent edges of the newly visited vertex to pq if they have not already been visited.

Finally, the function returns the MST set.

Example

Let's consider the following weighted undirected graph:

     1
  [0]---[1]
  / |   / |
3   | 2   |
  \ | /   |
   [3]---[2]
     6

The corresponding adjacency matrix is:

graph = [
    [0, 1, 3, 0],
    [1, 0, 2, 2],
    [3, 2, 0, 6],
    [0, 2, 6, 0]
]

Applying prim(graph) returns the following minimum spanning tree:

{(0, 1, 1), (1, 2, 2), (0, 3, 3)}

The minimum spanning tree contains the edges (0, 1), (1, 2), and (0, 3), with a total weight of 1+2+3=6.

Conclusion

Prim's algorithm is a simple and efficient algorithm for finding a minimum spanning tree in a weighted undirected graph. Using an adjacency matrix representation of the graph, the algorithm can be implemented using a priority queue and a set of visited vertices.