Prim's Algorithm: MST Using Priority Queue

Key Insights

Prim’s algorithm builds a minimum spanning tree by greedily selecting the smallest edge connecting the growing tree to an unvisited vertex, making the priority queue the critical data structure for efficient edge selection.
The key optimization is storing vertices (not edges) in the priority queue with their minimum known connection cost to the current tree, reducing both space and time complexity.
Choose Prim’s over Kruskal’s for dense graphs where E approaches V², as the O((V + E) log V) complexity with adjacency lists outperforms Kruskal’s union-find approach in these scenarios.

Introduction to Minimum Spanning Trees

A minimum spanning tree (MST) is a subset of edges from a connected, weighted, undirected graph that connects all vertices with the minimum possible total edge weight—without forming any cycles. If your graph has V vertices, the MST will always contain exactly V-1 edges.

MSTs solve real problems. When you’re laying network cables between offices, you want minimum total cable length while ensuring every office is connected. When clustering data points, MST-based algorithms group nearby points efficiently. Approximation algorithms for NP-hard problems like the Traveling Salesman Problem use MSTs as a starting point.

Two algorithms dominate MST construction: Kruskal’s and Prim’s. Kruskal’s sorts all edges globally and adds them if they don’t create cycles. Prim’s grows a single tree from a starting vertex, always adding the cheapest edge that expands the tree. Both are greedy and both produce optimal results, but their performance characteristics differ based on graph density.

How Prim’s Algorithm Works

Prim’s algorithm follows a simple greedy strategy:

Start with an arbitrary vertex as your initial tree
Find the minimum-weight edge connecting any tree vertex to any non-tree vertex
Add that edge and the new vertex to your tree
Repeat until all vertices are in the tree

Consider a graph with vertices A, B, C, D and edges: A-B (weight 1), A-C (weight 3), B-C (weight 2), B-D (weight 4), C-D (weight 1).

Starting from A:

Tree: {A}. Cheapest edge out: A-B (1). Add B.
Tree: {A, B}. Options: A-C (3), B-C (2), B-D (4). Cheapest: B-C (2). Add C.
Tree: {A, B, C}. Options: A-C (skip, C already in), B-D (4), C-D (1). Cheapest: C-D (1). Add D.
Tree: {A, B, C, D}. Done.

MST edges: A-B, B-C, C-D. Total weight: 4.

Here’s how we represent weighted graphs for Prim’s algorithm:

from collections import defaultdict

def create_graph():
    """Create adjacency list representation of weighted undirected graph."""
    graph = defaultdict(list)
    
    edges = [
        ('A', 'B', 1),
        ('A', 'C', 3),
        ('B', 'C', 2),
        ('B', 'D', 4),
        ('C', 'D', 1),
    ]
    
    for u, v, weight in edges:
        graph[u].append((v, weight))
        graph[v].append((u, weight))
    
    return graph

The adjacency list maps each vertex to a list of (neighbor, weight) tuples. For undirected graphs, add both directions.

The Role of the Priority Queue

The naive approach—scanning all edges to find the minimum at each step—gives O(V × E) complexity. For dense graphs, that’s O(V³). Unacceptable.

A priority queue (min-heap) reduces this to O(log V) per extraction. But here’s the insight that trips up many implementations: store vertices, not edges.

Each vertex outside the current tree has some minimum-cost edge connecting it to the tree (or infinity if no edge exists yet). We maintain these costs in the priority queue. When we add a vertex to the tree, we check if its neighbors now have cheaper connections and update accordingly.

# Priority queue operations for Prim's
# Entry format: (cost_to_connect, vertex)

import heapq

# Initialize: start vertex has cost 0, all others infinity
pq = [(0, start_vertex)]
min_cost = {v: float('inf') for v in vertices}
min_cost[start_vertex] = 0

# Extract minimum
cost, vertex = heapq.heappop(pq)

# Update neighbor's cost if we found a cheaper connection
if edge_weight < min_cost[neighbor]:
    min_cost[neighbor] = edge_weight
    heapq.heappush(pq, (edge_weight, neighbor))

One caveat: Python’s heapq doesn’t support decrease-key operations. We handle this by allowing duplicate entries and skipping vertices we’ve already processed. This adds some overhead but keeps the implementation simple.

Implementation with Min-Heap

Here’s the complete, production-ready implementation:

import heapq
from collections import defaultdict
from typing import Dict, List, Tuple, Set, Optional

def prim_mst(graph: Dict[str, List[Tuple[str, int]]], 
             start: str) -> Tuple[List[Tuple[str, str, int]], int]:
    """
    Compute MST using Prim's algorithm with priority queue.
    
    Args:
        graph: Adjacency list where graph[u] = [(v, weight), ...]
        start: Starting vertex
    
    Returns:
        Tuple of (list of MST edges as (u, v, weight), total weight)
    """
    if start not in graph:
        raise ValueError(f"Start vertex {start} not in graph")
    
    mst_edges: List[Tuple[str, str, int]] = []
    total_weight = 0
    in_mst: Set[str] = set()
    
    # Priority queue: (edge_weight, current_vertex, parent_vertex)
    # Parent is None for start vertex
    pq: List[Tuple[int, str, Optional[str]]] = [(0, start, None)]
    
    while pq and len(in_mst) < len(graph):
        weight, vertex, parent = heapq.heappop(pq)
        
        # Skip if already in MST (handles duplicate entries)
        if vertex in in_mst:
            continue
        
        # Add vertex to MST
        in_mst.add(vertex)
        
        # Record edge (skip for start vertex)
        if parent is not None:
            mst_edges.append((parent, vertex, weight))
            total_weight += weight
        
        # Add edges to unvisited neighbors
        for neighbor, edge_weight in graph[vertex]:
            if neighbor not in in_mst:
                heapq.heappush(pq, (edge_weight, neighbor, vertex))
    
    return mst_edges, total_weight


# Usage example
graph = defaultdict(list)
edges = [
    ('A', 'B', 1), ('A', 'C', 3), ('B', 'C', 2),
    ('B', 'D', 4), ('C', 'D', 1), ('D', 'E', 5),
]
for u, v, w in edges:
    graph[u].append((v, w))
    graph[v].append((u, w))

mst_edges, total = prim_mst(graph, 'A')
print(f"MST edges: {mst_edges}")
print(f"Total weight: {total}")

The key implementation details:

Track parent vertices: We store (weight, vertex, parent) tuples so we can reconstruct which edges form the MST.
Skip processed vertices: Since we can’t decrease keys, we push duplicates and skip vertices already in the MST.
Early termination: Stop when all vertices are included.

Time Complexity Analysis

Let’s break down the complexity with a binary heap:

Initialization: O(1) for the starting vertex
Main loop iterations: Each vertex is added to the MST exactly once: O(V)
Heap operations per vertex: We might push each edge once, and each pop is O(log V)
Total heap operations: O(E) pushes and O(E) pops in the worst case, each O(log V)

Final complexity: O((V + E) log V), which simplifies to O(E log V) for connected graphs.

For dense graphs where E ≈ V², this becomes O(V² log V). A Fibonacci heap can improve this to O(E + V log V) by supporting O(1) amortized decrease-key operations, but the constant factors make it rarely worthwhile in practice unless you’re dealing with massive graphs.

Practical Considerations and Edge Cases

Real-world graphs aren’t always clean. Here’s how to handle common issues:

def prim_mst_robust(graph: Dict[str, List[Tuple[str, int]]]) -> Tuple[List[Tuple[str, str, int]], int, bool]:
    """
    Robust Prim's implementation handling disconnected graphs.
    
    Returns:
        (mst_edges, total_weight, is_connected)
    """
    if not graph:
        return [], 0, True
    
    all_vertices = set(graph.keys())
    # Also include vertices that only appear as neighbors
    for neighbors in graph.values():
        for v, _ in neighbors:
            all_vertices.add(v)
    
    mst_edges = []
    total_weight = 0
    in_mst: Set[str] = set()
    
    # Process each connected component
    for start in all_vertices:
        if start in in_mst:
            continue
        
        pq = [(0, start, None)]
        
        while pq:
            weight, vertex, parent = heapq.heappop(pq)
            
            if vertex in in_mst:
                continue
            
            in_mst.add(vertex)
            
            if parent is not None:
                mst_edges.append((parent, vertex, weight))
                total_weight += weight
            
            for neighbor, edge_weight in graph.get(vertex, []):
                if neighbor not in in_mst:
                    heapq.heappush(pq, (edge_weight, neighbor, vertex))
    
    is_connected = len(mst_edges) == len(all_vertices) - 1
    return mst_edges, total_weight, is_connected

When to choose Prim’s over Kruskal’s:

Dense graphs (E close to V²): Prim’s with adjacency matrix is O(V²), while Kruskal’s sorting is O(E log E) = O(V² log V)
When you need to grow from a specific starting point
When the graph is already in adjacency list form

When to choose Kruskal’s:

Sparse graphs: Kruskal’s O(E log E) beats Prim’s when E « V²
When edges are already sorted or easy to sort
When using union-find is natural for your problem

Real-World Applications

MSTs appear everywhere in infrastructure optimization. Here’s a practical example:

def optimize_network_layout(offices: List[str], 
                           connection_costs: List[Tuple[str, str, int]]) -> dict:
    """
    Find minimum-cost network connecting all offices.
    
    Args:
        offices: List of office identifiers
        connection_costs: List of (office1, office2, cable_cost) tuples
    
    Returns:
        Dictionary with optimal connections and total cost
    """
    graph = defaultdict(list)
    for o1, o2, cost in connection_costs:
        graph[o1].append((o2, cost))
        graph[o2].append((o1, cost))
    
    # Ensure all offices are in graph
    for office in offices:
        if office not in graph:
            graph[office] = []
    
    mst_edges, total_cost, is_connected = prim_mst_robust(graph)
    
    if not is_connected:
        raise ValueError("Cannot connect all offices - network is disconnected")
    
    return {
        'connections': [(u, v) for u, v, _ in mst_edges],
        'total_cost': total_cost,
        'cost_breakdown': {f"{u}-{v}": w for u, v, w in mst_edges}
    }


# Example: Connect 5 offices with minimum cable cost
result = optimize_network_layout(
    offices=['HQ', 'NYC', 'LA', 'CHI', 'SEA'],
    connection_costs=[
        ('HQ', 'NYC', 100), ('HQ', 'LA', 200), ('HQ', 'CHI', 150),
        ('NYC', 'CHI', 80), ('LA', 'SEA', 120), ('CHI', 'SEA', 180),
        ('NYC', 'LA', 250), ('HQ', 'SEA', 300),
    ]
)
print(f"Install cables: {result['connections']}")
print(f"Total cost: ${result['total_cost']}")

Beyond network design, MSTs power image segmentation (treating pixels as graph vertices weighted by color similarity), phylogenetic tree construction in bioinformatics, and cluster analysis where cutting the longest MST edges produces natural groupings.

The priority queue optimization transforms Prim’s from a theoretical curiosity into a practical tool. Master it, and you’ll have a reliable solution for any minimum-cost connectivity problem you encounter.