Minimum Spanning Trees & Union-Find

A district wants to connect all its villages with optical fibre. Laying cable between some pairs of villages is cheap, between others it is expensive, and some pairs cannot be connected directly at all. Which cables should it lay so that every village is connected, directly or through others, at the lowest total cost? This is the minimum spanning tree problem.

In this lesson you will learn a small but powerful data structure, union-find, then use it in Kruskal's algorithm, compare it with Prim's algorithm, and finish with an MST trick for grouping similar items.

Spanning trees

A spanning tree of a connected, undirected graph is a set of edges that connects all V vertices with no cycles. Being a tree, it always has exactly V - 1 edges (see Trees & Binary Trees). A minimum spanning tree (MST) is a spanning tree whose total edge weight is as small as possible.

Why no cycles? If the chosen cables form a loop, you can remove the most expensive cable in that loop and every village stays connected, so a cheapest network never contains a loop.

Note the difference from shortest paths: an MST minimises the total cost of the whole network, not the distance between any particular pair of vertices.

Union-find

Both MST algorithms need to answer one question quickly: are these two vertices already connected by the edges chosen so far? Union-find (also called a disjoint set union, or DSU) keeps vertices in groups and supports two operations:

• find(x): return the representative (the root) of x's group. Two vertices are in the same group exactly when their roots are equal.
• union(a, b): merge the groups of a and b.

Each group is stored as a tree of parent links, where the root points to itself. Two simple tricks keep these trees very flat:

• Path compression: after find(x) walks up to the root, point every vertex on the way directly at the root.
• Union by size: attach the smaller tree under the root of the larger one.

With both tricks, a sequence of operations costs so close to O(1) each that, for any input you could ever store, it behaves as constant time. (Formally it is O(α(n)), where α is the inverse Ackermann function, which is at most 4 for all practical n.)

class UnionFind:
    def __init__(self, items):
        self.parent = {x: x for x in items}
        self.size = {x: 1 for x in items}

    def find(self, x):
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:              # path compression
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False                           # already together
        if self.size[ra] < self.size[rb]:          # union by size
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True

people = ["Aarav", "Priya", "Kabir", "Meera", "Rohan", "Isha"]
uf = UnionFind(people)
for a, b in [("Aarav", "Priya"), ("Kabir", "Meera"), ("Priya", "Kabir"), ("Rohan", "Isha")]:
    uf.union(a, b)

print(uf.find("Aarav") == uf.find("Meera"))
print(uf.find("Aarav") == uf.find("Isha"))
print(uf.union("Meera", "Aarav"))              # same group already

groups = {}
for p in people:
    groups.setdefault(uf.find(p), []).append(p)
print(list(groups.values()))

True
False
False
[['Aarav', 'Priya', 'Kabir', 'Meera'], ['Rohan', 'Isha']]

union returns False when the two are already in one group. That return value is exactly the "would this edge create a cycle?" test Kruskal's algorithm needs. Union-find is also the go-to structure for any "merge groups and ask if two things are connected" problem, such as friend circles or network connectivity, especially when connections keep arriving over time.

Kruskal's algorithm

Kruskal's algorithm is greedy (see Greedy Algorithms): it always takes the cheapest edge that is still useful.

1. Sort all edges by weight, cheapest first.
2. For each edge, if its two endpoints are in different groups, add it to the MST and union the groups. If they are already connected, skip it, because it would make a cycle.
3. Stop when the MST has V - 1 edges.

class UnionFind:
    def __init__(self, items):
        self.parent = {x: x for x in items}
    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]   # path halving
            x = self.parent[x]
        return x
    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        self.parent[rb] = ra
        return True

# Cost of laying fibre between villages, in Rs lakh (example figures)
cables = [("Badli", "Kharkhoda", 7), ("Badli", "Sampla", 5), ("Kharkhoda", "Sampla", 8),
          ("Kharkhoda", "Sonipat", 9), ("Sampla", "Rohtak", 7), ("Sampla", "Sonipat", 15),
          ("Rohtak", "Sonipat", 11), ("Rohtak", "Meham", 8), ("Sonipat", "Meham", 9),
          ("Sonipat", "Gohana", 6), ("Meham", "Gohana", 6)]
villages = {v for a, b, _ in cables for v in (a, b)}

def kruskal(vertices, edges):
    uf = UnionFind(vertices)
    chosen, total = [], 0
    for a, b, cost in sorted(edges, key=lambda e: e[2]):
        if uf.union(a, b):
            chosen.append((a, b, cost))
            total += cost
            if len(chosen) == len(vertices) - 1:
                break
    return chosen, total

mst, total = kruskal(villages, cables)
for a, b, cost in mst:
    print(f"{a:9} - {b:9} Rs {cost} lakh")
print("Total:", total, "lakh for", len(villages), "villages")

Badli     - Sampla    Rs 5 lakh
Sonipat   - Gohana    Rs 6 lakh
Meham     - Gohana    Rs 6 lakh
Badli     - Kharkhoda Rs 7 lakh
Sampla    - Rohtak    Rs 7 lakh
Rohtak    - Meham     Rs 8 lakh
Total: 39 lakh for 7 villages

Seven villages need exactly six cables. The cable Kharkhoda to Sampla (8) was skipped because those two were already connected through Badli. Sorting dominates the cost, so Kruskal runs in O(E log E). This version uses path halving, a one-line variant of path compression that points every other node at its grandparent.

Prim's algorithm

Prim's algorithm grows one tree outwards from any start vertex, always adding the cheapest edge that connects the tree to a vertex not yet in it. A min-heap holds the candidate edges, just like Dijkstra's algorithm in Shortest Paths: Dijkstra & Bellman-Ford. The difference is the key: Prim uses the weight of a single edge, Dijkstra uses the total distance from the source.

import heapq

cables = [("Badli", "Kharkhoda", 7), ("Badli", "Sampla", 5), ("Kharkhoda", "Sampla", 8),
          ("Kharkhoda", "Sonipat", 9), ("Sampla", "Rohtak", 7), ("Sampla", "Sonipat", 15),
          ("Rohtak", "Sonipat", 11), ("Rohtak", "Meham", 8), ("Sonipat", "Meham", 9),
          ("Sonipat", "Gohana", 6), ("Meham", "Gohana", 6)]
graph = {}
for a, b, c in cables:
    graph.setdefault(a, []).append((c, b))
    graph.setdefault(b, []).append((c, a))

def prim(graph, start):
    in_tree = {start}
    heap = [(cost, start, nxt) for cost, nxt in graph[start]]
    heapq.heapify(heap)
    chosen, total = [], 0
    while heap and len(in_tree) < len(graph):
        cost, a, b = heapq.heappop(heap)
        if b in in_tree:
            continue                         # both ends already in the tree
        in_tree.add(b)
        chosen.append((a, b, cost))
        total += cost
        for c, nxt in graph[b]:
            if nxt not in in_tree:
                heapq.heappush(heap, (c, b, nxt))
    return chosen, total

mst, total = prim(graph, "Rohtak")
print([f"{a}-{b}" for a, b, _ in mst])
print("Total:", total, "lakh")

['Rohtak-Sampla', 'Sampla-Badli', 'Badli-Kharkhoda', 'Rohtak-Meham', 'Meham-Gohana', 'Gohana-Sonipat']
Total: 39 lakh

Prim chose the same six cables in a different order, starting from Rohtak, and reached the same total of ₹39 lakh. With a binary heap it runs in O(E log V). Both algorithms are correct; Kruskal is usually simpler when you already have an edge list, and Prim suits graphs stored as adjacency lists, especially dense ones. If several edges share the same weight, there can be more than one MST, but every MST has the same total weight.

Advanced: clustering with an MST

Kruskal's algorithm joins the closest things first. If you stop early, when exactly k groups remain, you have split the vertices into k clusters where the groups are as far apart as possible. Here, delivery points in a city are grouped into 3 zones based on travel time between them:

# (continuing with the UnionFind class from the Kruskal example)
def clusters(points, links, k):
    uf = UnionFind(points)
    groups = len(points)
    for a, b, _ in sorted(links, key=lambda e: e[2]):
        if groups == k:
            break
        if uf.union(a, b):
            groups -= 1
    out = {}
    for p in points:
        out.setdefault(uf.find(p), []).append(p)
    return sorted(sorted(g) for g in out.values())

The full runnable version is in the exercise below. This "single-linkage clustering" is one of the simplest clustering methods used in data analysis.

Common mistakes

• Using an MST to find shortest paths. The MST minimises the total network cost; the path between two vertices inside it can be much longer than the shortest path.
• Forgetting the cycle check in Kruskal. Without union-find, a cheap edge between two already connected vertices sneaks in.
• Union-find without path compression or union by size, which can make each find O(n) on long chains.
• Running MST algorithms on a disconnected graph and expecting V - 1 edges. You get a spanning forest instead; check the edge count.

Key points

• A spanning tree connects all V vertices with V - 1 edges and no cycles; an MST has the minimum total weight.
• Union-find supports find and union in nearly O(1) amortised time with path compression and union by size.
• Kruskal sorts edges and adds each one that joins two different groups: O(E log E).
• Prim grows one tree from a start vertex using a min-heap of edges: O(E log V).
• Every MST of a graph has the same total weight; stopping Kruskal at k groups gives k clusters.

class UnionFind:
    def __init__(self, items):
        self.parent = {x: x for x in items}
    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x
    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        self.parent[rb] = ra
        return True

def clusters(points, links, k):
    uf = UnionFind(points)
    groups = len(points)
    for a, b, _ in sorted(links, key=lambda e: e[2]):
        if groups == k:
            break
        if uf.union(a, b):
            groups -= 1
    out = {}
    for p in points:
        out.setdefault(uf.find(p), []).append(p)
    return sorted(sorted(g) for g in out.values())

links = [("Kothrud", "Karve Nagar", 6), ("Kothrud", "Aundh", 18), ("Karve Nagar", "Deccan", 9),
         ("Deccan", "Shivajinagar", 5), ("Shivajinagar", "Aundh", 10), ("Aundh", "Baner", 7),
         ("Baner", "Wakad", 8), ("Deccan", "Camp", 14), ("Shivajinagar", "Camp", 12)]
points = sorted({p for a, b, _ in links for p in (a, b)})
print(clusters(points, links, 3))
print(clusters(points, links, 2))

[['Aundh', 'Baner', 'Wakad'], ['Camp'], ['Deccan', 'Karve Nagar', 'Kothrud', 'Shivajinagar']]
[['Aundh', 'Baner', 'Deccan', 'Karve Nagar', 'Kothrud', 'Shivajinagar', 'Wakad'], ['Camp']]