Skip to main content

​
Overview

Prim’s algorithm finds a minimum spanning tree (MST) for a weighted undirected graph. An MST is a subset of edges that connects all vertices with the minimum total edge weight, without forming any cycles.
This implementation works with both AdjacencyListWeightedGraph and AdjacencyMatrixWeightedGraph from the DSA package.

​
Functions

​
prims_mst

Computes the minimum spanning tree of a graph.
prim.py:5
def prims_mst(graph, start: str, mst_graph=None) -> AdjacencyListWeightedGraph:
    """
    Returns an MST given a graph and starting vertex.
    (Future: return a Tree type instead of a Graph type)

    Args:
        graph: The graph to search an MST from. (can be either an AdjacencyListWeightedGraph or AdjacencyMatrixWeightedGraph)
        start (string): The starting vertex label.
        mst_graph: An empty graph object to output the MST in to.

    Returns:
        AdjacencyListWeightedGraph: the MST of the graph.
    """
Parameters:
  • graph: The weighted graph to find an MST from (supports both adjacency list and matrix representations)
  • start (str): The starting vertex label
  • mst_graph (optional): An empty graph object to populate with the MST. If not provided, a new AdjacencyListWeightedGraph is created
Returns:
  • AdjacencyListWeightedGraph: The minimum spanning tree of the input graph

​
mst_weight

Calculates the total weight of a graph (typically an MST).
prim.py:46
def mst_weight(graph) -> int:
    """
    Returns the total weight of a graph given a starting vertex
    
    Args:
        graph: The graph to find the total edge weight of.

    Returns:
        int: The total weight of the graph.
    """
Parameters:
  • graph: The graph to calculate the total weight for
Returns:
  • int: The sum of all edge weights in the graph

​
Usage Examples

1

Create a weighted graph

First, create a weighted graph with vertices and edges:
from dsa.graph import AdjacencyListWeightedGraph
from dsa.prim import prims_mst, mst_weight

# Create a weighted graph
graph = AdjacencyListWeightedGraph()
graph.add_edge('A', 'B', 4)
graph.add_edge('A', 'C', 2)
graph.add_edge('B', 'C', 1)
graph.add_edge('B', 'D', 5)
graph.add_edge('C', 'D', 8)
graph.add_edge('C', 'E', 10)
graph.add_edge('D', 'E', 2)
2

Find the minimum spanning tree

Use prims_mst() to compute the MST:
# Find MST starting from vertex 'A'
mst = prims_mst(graph, 'A')

# Display the edges in the MST
for start, end, weight in mst.edges():
    print(f"{start} -- {end}: {weight}")
3

Calculate the total MST weight

Use mst_weight() to get the total weight of the MST:
# Calculate the total weight
total_weight = mst_weight(mst)
print(f"Total MST weight: {total_weight}")
4

Visualize the MST (optional)

Use the visualization tools to see the MST:
from dsa.draw import GraphDraw

# Draw the original graph with MST highlighted
gd = GraphDraw(graph)
gd.draw(show_mst=True)

# Or draw only the MST
gd.draw(mst_only=True)

​
Algorithm Details

The implementation uses:
  1. Priority Queue: Efficiently selects the minimum weight edge (prim.py:30)
  2. Visited Set: Tracks which vertices have been added to the MST (prim.py:29)
  3. Edge Selection: Only adds edges that connect to unvisited vertices (prim.py:40)
  4. Helper Function: add_adjacent() adds all edges from a vertex to the priority queue (prim.py:18-23)
def add_adjacent(graph, pq: PriorityQueue, visited: set, node: str):
    """Add all adjacent vertices from the given node to the priority queue."""
    visited.add(node)
    for adjacent, weight in graph[node].items():
        if adjacent not in visited:
            pq.push(weight, (node, adjacent))  # Push edge with weight as priority

​
Time Complexity

  • Time Complexity: O((V + E) log V) where V is the number of vertices and E is the number of edges
  • Space Complexity: O(V + E) for storing the MST and priority queue
The algorithm continues until all vertices are visited or the priority queue is empty (prim.py:36), ensuring that all reachable vertices are included in the MST.

​
Key Properties

  • The resulting MST is always a tree (connected and acyclic)
  • For a connected graph with V vertices, the MST contains exactly V-1 edges
  • The total weight of the MST is minimal among all possible spanning trees
  • The choice of starting vertex does not affect the total weight of the MST

Build docs developers (and LLMs) love

Get started for freeTalk to us