Program to find out an MST using Prim's algorithm in Python

A Minimum Spanning Tree (MST) is a subset of edges in a weighted graph that connects all vertices with the minimum possible total edge weight and no cycles. Prim's algorithm builds the MST by starting from a vertex and greedily adding the minimum weight edge that connects a visited vertex to an unvisited one.

Algorithm Steps

Prim's algorithm follows these steps ?

• Start with any vertex and mark it as visited
• Find the minimum weight edge connecting visited to unvisited vertices
• Add this edge to the MST and mark the new vertex as visited
• Repeat until all vertices are visited

Implementation

Here's the complete implementation using Prim's algorithm ?

def mst_find(G, s):
    distance = [float("inf")] * len(G)
    distance[s] = 0
    visited = [False] * len(G)
    total_cost = 0
    
    while True:
        min_weight = float("inf")
        min_idx = -1
        
        # Find minimum weight edge to unvisited vertex
        for i in range(len(G)):
            if not visited[i] and distance[i] < min_weight:
                min_weight = distance[i]
                min_idx = i
        
        # If no vertex found, MST is complete
        if min_idx == -1:
            break
            
        # Add vertex to MST
        total_cost += min_weight
        visited[min_idx] = True
        
        # Update distances to adjacent vertices
        for neighbor, weight in G[min_idx].items():
            if not visited[neighbor]:
                distance[neighbor] = min(distance[neighbor], weight)
    
    return total_cost

def solve(n, edges, start_vertex):
    # Create adjacency list representation
    graph = {i: {} for i in range(n)}
    
    for u, v, w in edges:
        # Convert to 0-based indexing
        u -= 1
        v -= 1
        
        # Handle multiple edges between same vertices (keep minimum)
        if v in graph[u]:
            graph[u][v] = min(graph[u][v], w)
            graph[v][u] = min(graph[v][u], w)
        else:
            graph[u][v] = w
            graph[v][u] = w
    
    return mst_find(graph, start_vertex)

# Example usage
n = 4
edges = [(1, 2, 5), (1, 3, 5), (2, 3, 7), (1, 4, 4)]
start_vertex = 3

result = solve(n, edges, start_vertex)
print(f"Minimum Spanning Tree cost: {result}")

Minimum Spanning Tree cost: 14

How It Works

The algorithm maintains two key data structures ?

• distance[] − Stores the minimum weight to reach each unvisited vertex
• visited[] − Tracks which vertices are already in the MST

Starting from vertex 3, it adds edges in this order: (3?1, weight=5), (1?4, weight=4), (1?2, weight=5), giving total cost 14.

Time Complexity

Conclusion

Prim's algorithm efficiently finds the MST by greedily selecting minimum weight edges. The implementation uses adjacency lists for graph representation and achieves O(V²) time complexity, making it suitable for dense graphs.