Prim's Algorithm

1import heapq
2from collections import defaultdict
3
4def prim(n: int, edges: list[tuple[int, int, int]]) -> list[tuple[int, int, int]]:
5    """Find Minimum Spanning Tree using Prim's algorithm.
6
7    Args:
8        n: Number of vertices (0 to n-1)
9        edges: List of (u, v, weight) tuples
10
11    Returns:
12        List of edges in the MST as (u, v, weight) tuples
13    """
14    # Build adjacency list: node -> [(neighbor, weight), ...]
15    graph: dict[int, list[tuple[int, int]]] = defaultdict(list)
16    for u, v, w in edges:
17        graph[u].append((v, w))
18        graph[v].append((u, w))
19
20    mst: list[tuple[int, int, int]] = []
21    visited: set[int] = set()
22
23    # Min-heap: (weight, from_node, to_node)
24    # Start from node 0
25    heap: list[tuple[int, int, int]] = [(0, -1, 0)]
26
27    while heap and len(visited) < n:
28        weight, from_node, to_node = heapq.heappop(heap)
29
30        if to_node in visited:
31            continue
32
33        visited.add(to_node)
34
35        # Add edge to MST (skip the initial dummy edge)
36        if from_node != -1:
37            mst.append((from_node, to_node, weight))
38
39        # Explore neighbors
40        for neighbor, edge_weight in graph[to_node]:
41            if neighbor not in visited:
42                heapq.heappush(heap, (edge_weight, to_node, neighbor))
43
44    return mst