Graph Algorithms: Shortest Paths & MST

Dijkstra's Algorithm

Starting from a source node, repeatedly expand whichever unvisited node currently has the smallest known distance, exactly Lesson 9's min-heap giving O(log n) access to "the best option so far." Expanding a node relaxes its neighbors: if reaching a neighbor through the node just expanded is cheaper than the best known route so far, update it.

import heapq

def dijkstra(graph: dict[str, list[tuple[str, int]]], source: str) -> dict[str, float]:
    distances = {node: float("inf") for node in graph}
    distances[source] = 0
    visited = set()
    heap = [(0, source)]
    while heap:
        dist, node = heapq.heappop(heap)      # always expand the closest unvisited node
        if node in visited:
            continue
        visited.add(node)
        for neighbor, weight in graph[node]:
            new_dist = dist + weight
            if new_dist < distances[neighbor]:
                distances[neighbor] = new_dist
                heapq.heappush(heap, (new_dist, neighbor))
    return distances

graph = {
    "A": [("B", 4), ("C", 1)],
    "B": [("D", 1)],
    "C": [("B", 2), ("D", 5)],
    "D": [],
}
print(dijkstra(graph, "A"))

{'A': 0, 'B': 3, 'C': 1, 'D': 4}

Here is that graph, with each edge's weight and the final shortest distance from A under every node:

Watch the distances tighten as each node is finalized. Every step pulls the closest unvisited node off the heap, locks its distance in (green), and relaxes its neighbors (amber). Notice B drops from 4 to 3 once the cheaper route through C is found, and D drops from 6 to 4:

Bellman-Ford: Handling Negative Weights

Instead of greedily finalizing one node at a time, Bellman-Ford relaxes every edge in the graph, repeated V - 1 times (a shortest path between any two nodes visits at most every other node once, so no path needs more relaxation rounds than that). It's slower than Dijkstra, but it stays correct no matter what the edge weights are, and one extra pass over all the edges doubles as a negative cycle detector: if any edge can still be relaxed after V - 1 rounds, a negative cycle exists and no shortest path is even well-defined.

def bellman_ford(graph: dict[str, list[tuple[str, int]]], source: str) -> dict[str, float] | None:
    distances = {node: float("inf") for node in graph}
    distances[source] = 0
    edges = [(u, v, w) for u in graph for v, w in graph[u]]

    for _ in range(len(graph) - 1):          # a shortest path visits at most V-1 edges
        for u, v, w in edges:
            if distances[u] + w < distances[v]:
                distances[v] = distances[u] + w

    for u, v, w in edges:                     # one more pass: still improving means a negative cycle
        if distances[u] + w < distances[v]:
            return None
    return distances

Here's a graph where Dijkstra's greedy shortcut actually produces the wrong answer. A reaches C both directly (cost 1) and via B (cost 2, then -5, netting -3), and the cheap direct route gets C finalized before Dijkstra ever discovers the better path through B:

graph = {
    "A": [("C", 1), ("B", 2)],
    "B": [("C", -5)],
    "C": [("D", 1)],
    "D": [],
}
print("Dijkstra:     ", dijkstra(graph, "A"))
print("Bellman-Ford: ", bellman_ford(graph, "A"))

Dijkstra:      {'A': 0, 'B': 2, 'C': -3, 'D': 2}
Bellman-Ford:  {'A': 0, 'B': 2, 'C': -3, 'D': -2}

The true shortest A-to-D distance is 2 + (-5) + 1 = -2, exactly what Bellman-Ford finds. Dijkstra reports 2, it already committed to reaching D through the direct A→C edge before the cheaper route through B was even discovered, and it never goes back to reconsider, the same kind of irrevocable commitment that broke greedy coin change back in Lesson 17.

A Minimum Spanning Tree Needs a Cycle Check

A minimum spanning tree (MST) is the cheapest set of edges that connects every node in a graph, using exactly V - 1 edges and containing no cycles. Building one greedily by adding cheap edges needs a fast way to answer one question repeatedly: would adding this edge create a cycle, meaning its two endpoints are already connected through edges already chosen? The Union-Find (disjoint-set) structure answers that in nearly O(1) time:

class UnionFind:
    def __init__(self, nodes: list[str]) -> None:
        self.parent = {n: n for n in nodes}
        self.rank = {n: 0 for n in nodes}

    def find(self, x: str) -> str:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])   # path compression
        return self.parent[x]

    def union(self, x: str, y: str) -> bool:
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False              # already connected, this edge would form a cycle
        if self.rank[root_x] < self.rank[root_y]:
            root_x, root_y = root_y, root_x
        self.parent[root_y] = root_x
        if self.rank[root_x] == self.rank[root_y]:
            self.rank[root_x] += 1
        return True

Each node starts as its own isolated group. find walks up to a group's representative, compressing the path along the way so future lookups are faster. union merges two groups (attaching the smaller-ranked tree under the larger, keeping lookups shallow) and returns False if the two nodes were already in the same group, exactly the cycle check needed.

Kruskal's Algorithm

Sort every edge by weight, then greedily add each one, cheapest first, unless it would connect two nodes already in the same group (a cycle). Union-Find makes that check fast enough for the whole approach to be dominated by the initial sort:

def kruskal(nodes: list[str], edges: list[tuple[int, str, str]]) -> tuple[list[tuple[int, str, str]], int]:
    uf = UnionFind(nodes)
    mst_edges = []
    total_weight = 0
    for weight, u, v in sorted(edges):        # cheapest edges first
        if uf.union(u, v):                    # only keep it if it doesn't close a cycle
            mst_edges.append((weight, u, v))
            total_weight += weight
    return mst_edges, total_weight

nodes = ["A", "B", "C", "D", "E"]
edges = [
    (2, "A", "B"), (3, "A", "C"), (3, "B", "C"),
    (4, "B", "D"), (5, "C", "D"), (1, "C", "E"), (6, "D", "E"),
]
mst_edges, total_weight = kruskal(nodes, edges)
print(mst_edges)
print(total_weight)

[(1, 'C', 'E'), (2, 'A', 'B'), (3, 'A', 'C'), (4, 'B', 'D')]
10

Prim's Algorithm

Prim's takes a different angle on the same problem: grow a single tree outward from a starting node, at each step adding the cheapest edge that connects a node already in the tree to one that isn't. Structurally, it's Dijkstra's algorithm with one change, track the cost of the connecting edge instead of the cumulative distance from the source:

import heapq

def prim(graph: dict[str, list[tuple[str, int]]], start: str) -> tuple[list[tuple[int, str, str]], int]:
    visited = {start}
    mst_edges = []
    total_weight = 0
    heap = [(weight, start, neighbor) for neighbor, weight in graph[start]]
    heapq.heapify(heap)
    while heap and len(visited) < len(graph):
        weight, u, v = heapq.heappop(heap)    # cheapest edge crossing the frontier
        if v in visited:
            continue
        visited.add(v)
        mst_edges.append((weight, u, v))
        total_weight += weight
        for neighbor, w in graph[v]:
            if neighbor not in visited:
                heapq.heappush(heap, (w, v, neighbor))
    return mst_edges, total_weight

# same nodes and edges as the Kruskal's example above
nodes = ["A", "B", "C", "D", "E"]
edges = [
    (2, "A", "B"), (3, "A", "C"), (3, "B", "C"),
    (4, "B", "D"), (5, "C", "D"), (1, "C", "E"), (6, "D", "E"),
]
graph = {n: [] for n in nodes}
for w, u, v in edges:
    graph[u].append((v, w))
    graph[v].append((u, w))

mst_edges, total_weight = prim(graph, "A")
print(mst_edges)
print(total_weight)

[(2, 'A', 'B'), (3, 'A', 'C'), (1, 'C', 'E'), (4, 'B', 'D')]
10

Same graph, same total weight (10) as Kruskal's, just a different edge order and no need for Union-Find, Prim's already knows every edge it considers touches the growing tree. Both algorithms are correct; which one's faster depends on how dense the graph is, covered in the complexity table below.

Complexity at a Glance

E is at most V², so O(E log E) and O(E log V) describe the same growth rate, meaning the binary-heap versions of Kruskal's and Prim's shown here are actually in the same complexity class. Prim's classic advantage on very dense graphs comes from a different implementation, an adjacency matrix with no heap at all, giving O(V²), better than O(E log V) once E gets close to V², but that trade only pays off in that dense-graph regime. Bellman-Ford's O(V × E) is strictly worse than Dijkstra's, the price paid for correctness under negative weights.