Union-Find and Minimum Spanning Trees

A power company must join six towns with cable. Each possible cable has a price, and the company wants every town connected for the smallest total price. Before it can even start, it needs to answer a simpler question again and again: are these two towns already joined? This lesson builds a tiny structure that answers that question almost instantly, then uses it to find the cheapest network.

Learning Objectives

After this lesson, you will be able to:

Explain the parent array and trace find and union by hand on a small example
Say what path compression and union by rank each do, and show with counted steps why the two together are almost O(1) per operation
State in plain words what the inverse Ackermann function is, and where the formal amortized argument lives
Use Union-Find for a stream of new links, for a cycle in an undirected graph, for grid islands and for merging accounts
Say what a spanning tree is, state the cut property in plain words, and say why it makes greedy choices safe
Write Kruskal with sort and Union-Find, write Prim with heapq, give the time and space of each, and choose between them
Explain why the weight of a minimum spanning tree is unique, when the tree itself is unique, and how single-linkage clustering is Kruskal in disguise

Prerequisites

Before You Start

Review

#Two Questions About Connections

Think of islands joined by bridges. Someone asks: can I walk from island A to island B? Then a new bridge opens, and someone asks again. A graph search such as breadth-first search answers each question by walking the whole map. That is fine once. It is slow if the map changes and the questions never stop.

In this lesson, a graph is a set of nodes (also called vertices) joined by edges. A group is a set of nodes that can reach each other. Programmers also say connected component. The first question is: which group is this node in, and are these two nodes in the same group?

The second question is about cost. Each edge has a weight, such as a cable price. Which edges join every node for the smallest total weight? That cheapest joined network is a minimum spanning tree, and we build it in the second half.

#Quick Recap: Graphs and Heaps

Five lines of what you need.

A graph has V nodes and E edges. An edge can carry a weight.
An edge list is a plain list of (weight, u, v) triples. An adjacency list gives each node its neighbours (Graphs).
A cycle is a path that leaves a node and returns to it without reusing an edge.
A heap hands you the smallest value in O(log n), and heapq is Python's version (Heaps).
sorted costs O(n log n) time and O(n) extra space (Sorting Algorithms).

#The Parent Array Idea

Here is the whole structure. Give every item a number from 0 to n-1. Keep a list called parent. Item i has a parent link, a pointer that is just another item's number, stored in parent[i]. At the start nobody has a real parent, so each item points at itself: parent = [0, 1, 2, 3, 4, 5]. An item that points at itself is a root. A root names its group.

To find the group of an item, follow the parent links upward until you reach a root. That operation is called find. Two items are in the same group exactly when their roots are the same item.

To join two groups, find both roots and hang one root under the other. That operation is called union. We only ever change the parent of a root. Trace it by hand first.

Step	Call	parent after the call	Groups
start		`[0, 1, 2, 3, 4, 5]`	six groups of one
1	`union(0, 1)`	`[0, 0, 2, 3, 4, 5]`	1 goes under 0
2	`union(2, 3)`	`[0, 0, 2, 2, 4, 5]`	3 goes under 2
3	`union(1, 3)`	`[0, 0, 0, 2, 4, 5]`	root of 1 is 0, root of 3 is 2, so 2 goes under 0

Step 3 is the one to study. We did not touch items 1 and 3 themselves. We found their roots, 0 and 2, and linked the roots. Here is the picture of the forest after step 3. A forest is a set of trees, one per group.

text

   0       4     5
  / \
 1   2
     |
     3

Now the code. It has the three pieces from the table: a list, a find that climbs, and a union that links two roots.

pythonrunnable cell

parent = list(range(6))          # every item starts as its own root


def find(x):
    """Walk up the parent links until we reach a root."""
    while parent[x] != x:
        x = parent[x]
    return x


def union(a, b):
    """Join the groups of a and b by hanging one root under the other."""
    root_a, root_b = find(a), find(b)
    if root_a != root_b:
        parent[root_b] = root_a


union(0, 1)
print(parent)
union(2, 3)
print(parent)
union(1, 3)
print(parent)
print(find(3), find(4))
print(find(1) == find(3), find(1) == find(4))
# Output:
# [0, 0, 2, 3, 4, 5]
# [0, 0, 2, 2, 4, 5]
# [0, 0, 0, 2, 4, 5]
# 0 4
# True False

The printed lists match the table row by row. Item 3 climbs 3, then 2, then 0, so its root is 0. Item 4 never moved, so it is its own root.

What Do You Think?

Start from parent = [0, 0, 2, 2, 4, 5] (the state after step 2 above). The code is the plain version: it hangs the second root under the first. What is parent after union(3, 5)?

#Why the Plain Version Gets Slow

Each find costs as many steps as the tree is tall. Nothing in the plain code keeps trees short. Join item 1 to item 0, then 2 to that group, then 3, and so on, always hanging the old group under the new item. The group becomes one long chain, and find(0) must climb the whole chain.

pythonrunnable cell

def chain_hops(n):
    """Plain union only. Join item i+1 to the group of item i, for every i."""
    parent = list(range(n))

    def find(x):
        hops = 0
        while parent[x] != x:
            x = parent[x]
            hops += 1
        return x, hops

    for i in range(n - 1):
        root_a, _ = find(i + 1)
        root_b, _ = find(i)
        if root_a != root_b:
            parent[root_b] = root_a
    return find(0)[1]


for n in (10, 100, 1_000):
    print(n, chain_hops(n))
# Output:
# 10 9
# 100 99
# 1000 999

One find took n-1 steps. So the plain structure costs O(n) per operation in the worst case, no better than a graph search. Two small tricks fix this, and each attacks a different cause of tall trees.

#Trick One: Path Compression

The first trick repairs the tree while you use it. After find reaches the root, it already knows the root. So it walks the path a second time and points every item on the way straight at the root. This is path compression. The next find on any of those items takes one step.

pythonrunnable cell

parent = [1, 2, 3, 4, 4]          # a chain: 0 -> 1 -> 2 -> 3 -> 4, and 4 is the root


def find(x):
    """Find the root, then point every item on the path straight at it."""
    root = x
    while parent[root] != root:       # pass 1: walk up to the root
        root = parent[root]
    while parent[x] != root:          # pass 2: re-point each item on the path
        parent[x], x = root, parent[x]
    return root


print(parent)
print(find(0))
print(parent)
# Output:
# [1, 2, 3, 4, 4]
# 4
# [4, 4, 4, 4, 4]

The first find(0) paid for the long climb, and in return the whole chain became flat. We wrote find with two loops and not with recursion, so a very tall tree cannot hit Python's recursion limit. Compression costs only a second walk over a path you just walked, so it at most doubles the cost of one find.

#Trick Two: Union by Rank

Compression helps after trees are tall. The second trick stops them growing tall in the first place. When you link two roots, hang the shorter tree under the taller one. The tall tree's height does not change, and the combined tree is not much taller than the larger part.

We do not want to measure heights exactly, because compression changes them. So each root keeps a number called its rank, an upper bound on its height. A new item has rank 0. When two roots of equal rank are linked, the winner's rank goes up by one. Otherwise no rank changes. Some books keep the size of each tree instead and hang the smaller under the larger. Both work equally well.

What Do You Think?

We build 4,096 items into one group by joining equal-size trees: items in pairs, then the pairs in pairs, and so on. Union by rank decides every link, and there is no path compression. How many steps does the deepest item need to reach its root?

pythonrunnable cell

def deepest(n, pairs):
    """Union by rank, no compression. Return how many steps the deepest item is from its root."""
    parent = list(range(n))
    rank = [0] * n

    def find(x):
        while parent[x] != x:
            x = parent[x]
        return x

    for a, b in pairs:
        ra, rb = find(a), find(b)
        if ra == rb:
            continue
        if rank[ra] < rank[rb]:
            ra, rb = rb, ra
        if rank[ra] == rank[rb]:
            rank[ra] += 1
        parent[rb] = ra
    depth = 0
    for x in range(n):
        steps, y = 0, x
        while parent[y] != y:
            y, steps = parent[y], steps + 1
        depth = max(depth, steps)
    return depth


def pairs_of_equal_trees(n):
    """Join items in pairs, then join the pairs, then join those, and so on."""
    pairs, step = [], 1
    while step < n:
        pairs += [(i, i + step) for i in range(0, n, 2 * step)]
        step *= 2
    return pairs


for n in (16, 256, 4096):
    print(n, deepest(n, pairs_of_equal_trees(n)))
# Output:
# 16 4
# 256 8
# 4096 12

The depths 4, 8 and 12 are log base 2 of 16, 256 and 4,096. That is no accident. A tree of height h can only exist if two trees of height h-1 were joined, so it holds at least twice as many items as a tree one level shorter. A tree with n items is therefore at most log n levels tall. Union by rank alone makes find cost O(log n) in the worst case.

#The Finished UnionFind Class

Put the pieces together. This is the class used across this part of the track: integer ids, a parent list, a rank list, a count of components, and the methods find, union and connected. union returns True when it really merged two groups, and False when they were already one. That return value powers most of the uses below.

pythonrunnable cell

class UnionFind:
    """Groups of items numbered 0 to n-1, with path compression and union by rank."""

    def __init__(self, n: int) -> None:
        """Start with n groups of one item each."""
        self.parent: list[int] = list(range(n))   # each item is its own root
        self.rank: list[int] = [0] * n            # upper bound on tree height
        self.components: int = n                  # how many groups there are now

    def find(self, x: int) -> int:
        """Return the root of x's group, flattening the path on the way."""
        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, x: int, y: int) -> bool:
        """Merge the groups of x and y. Return True if they were different groups."""
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        if self.rank[root_x] < self.rank[root_y]:  # union by rank:
            root_x, root_y = root_y, root_x        # the taller tree stays on top
        self.parent[root_y] = root_x
        if self.rank[root_x] == self.rank[root_y]:
            self.rank[root_x] += 1
        self.components -= 1
        return True

    def connected(self, x: int, y: int) -> bool:
        """Say whether x and y are in the same group."""
        return self.find(x) == self.find(y)


uf = UnionFind(5)
print(uf.parent, uf.components)
uf.union(0, 1)
print(uf.parent, uf.components)
uf.union(2, 3)
print(uf.parent, uf.components)
uf.union(1, 3)
print(uf.parent, uf.rank, uf.components)
print(uf.find(3), uf.parent)
print(uf.connected(1, 3), uf.connected(1, 4))
print(uf.union(0, 3))
# Output:
# [0, 1, 2, 3, 4] 5
# [0, 0, 2, 3, 4] 4
# [0, 0, 2, 2, 4] 3
# [0, 0, 0, 2, 4] [2, 0, 1, 0, 0] 2
# 0 [0, 0, 0, 0, 4]
# True False
# False

Read the fourth line. When we joined 1 and 3, their roots were 0 and 2. Both had rank 1, so the tie goes to the first, 0 stays on top, and its rank rises to 2. Then find(3) climbed 3, 2, 0 and pointed item 3 straight at 0. The last line shows union(0, 3) returning False because they were already together.

The costs, with reasons. Building the class is O(n) time and O(n) space, for two lists of n numbers. find and union take almost constant time per call, explained next. The extra space per operation is O(1), because both loops are iterative.

#How Fast Is It Really?

The claim is bold, so count steps. The next cell runs the same workloads under four setups: neither trick, only compression, only rank, and both. A hop is one step up a parent link. The chain workload is the bad case from before. The mixed workload links items in a scrambled but fixed pattern. Each run does all its unions, then one find per item.

pythonrunnable cell

def total_hops(n, pairs, compress, by_rank):
    """Do all the unions, then one find per item. Count every step up a parent link."""
    parent = list(range(n))
    rank = [0] * n
    hops = 0

    def find(x):
        nonlocal hops
        root = x
        while parent[root] != root:
            root = parent[root]
            hops += 1
        if compress:
            while parent[x] != root:
                parent[x], x = root, parent[x]
        return root

    for a, b in pairs:
        ra, rb = find(a), find(b)
        if ra == rb:
            continue
        if by_rank:
            if rank[ra] < rank[rb]:
                ra, rb = rb, ra
            if rank[ra] == rank[rb]:
                rank[ra] += 1
        parent[rb] = ra
    for x in range(n):
        find(x)
    return hops


n = 4096
chain = [(i + 1, i) for i in range(n - 1)]
mixed = [(i, (i * 7919 + 13) % n) for i in range(n)]
print("tricks            chain   mixed")
for label, compress, by_rank in (("neither", False, False),
                                 ("compression only", True, False),
                                 ("rank only", False, True),
                                 ("both", True, True)):
    print(f"{label:<16} {total_hops(n, chain, compress, by_rank):>7} {total_hops(n, mixed, compress, by_rank):>7}")

print()
print("both tricks, mixed workload")
for n in (1_000, 10_000, 100_000):
    pairs = [(i, (i * 7919 + 13) % n) for i in range(n)]
    hops = total_hops(n, pairs, True, True)
    print(n, hops, f"{hops / (3 * n):.2f} hops per find")
# Output:
# tricks            chain   mixed
# neither          8386560   26553
# compression only    8189   12019
# rank only           8188   18565
# both                8188   11983
#
# both tricks, mixed workload
# 1000 2426 0.81 hops per find
# 10000 27721 0.92 hops per find
# 100000 295450 0.98 hops per find

With neither trick the chain costs 8,386,560 hops, which is 4,096 times 4,095 divided by 2, the O(n²) pattern. Either trick brings it to about two hops per item. In the second table the number of items grows a hundred times, and the average stays below one hop per find. That flat line is what "almost constant" looks like.

Now the plain-words version of the famous bound.

With both tricks, a run of m operations on n items costs O((m + n) times α(n)) in total, where α is the inverse Ackermann function. The Ackermann function is a function that grows faster than any tower of powers you can write down. Its inverse is the opposite: it grows so slowly that it stays at 4 or below for any n that could ever be stored. So α(n) behaves like a small constant, and each operation is O(α(n)) amortized, which is almost O(1).

Amortized means averaged over the whole run, so one call may cost more while the average stays tiny. Robert Tarjan proved this bound in the 1970s. We will not fake a proof here. The formal argument uses the potential method, and the asymptotics and amortized analysis lesson builds the tools for it. The counted steps above are the evidence you can check today.

Quick check

You keep path compression but remove union by rank, so the root of the second group always goes under the root of the first. Which statement is true about the cost of an operation?

#Using Union-Find: Four Classic Jobs

All four jobs below use the same class. Each time, ask: what are the items, and what counts as a link?

#Groups That Change Over Time

Items are people, links are friendships. Friendships arrive one at a time, and after each we want the number of groups. A fresh breadth-first search after every link costs O(V + E) each time. With Union-Find, each link costs almost O(1).

pythonrunnable cell

uf = UnionFind(6)
for a, b in [(0, 1), (2, 3), (1, 3), (4, 5), (0, 2)]:
    merged = uf.union(a, b)
    print((a, b), "merged" if merged else "already together", uf.components)
# Output:
# (0, 1) merged 5
# (2, 3) merged 4
# (1, 3) merged 3
# (4, 5) merged 2
# (0, 2) already together 2

The last link joined two people who were already in one group, so the count stayed at 2. For m links over n people, the total is O((n + m) α(n)) time and O(n) space.

#A Cycle in an Undirected Graph

An undirected edge joins two nodes both ways. If you add an edge between two nodes that are already in one group, there was already a path between them, and the new edge closes a loop. So union returning False is exactly a cycle.

pythonrunnable cell

def has_cycle(n, edges):
    """True if the undirected graph has a cycle. O(n + E * alpha(n)) time, O(n) space."""
    uf = UnionFind(n)
    for u, v in edges:
        if not uf.union(u, v):       # both ends already joined, so this edge closes a loop
            return True
    return False


print(has_cycle(4, [(0, 1), (1, 2), (2, 3)]))
print(has_cycle(4, [(0, 1), (1, 2), (2, 3), (3, 1)]))
print(has_cycle(3, [(0, 1), (1, 1)]))
# Output:
# False
# True
# True

The third graph has an edge from node 1 to itself, a self-loop, and that counts as a cycle. Two copies of the same edge also count. This test is for undirected graphs only. In a directed graph, an edge A to B and an edge B to C do not mean you can come back, so the topological sort of the shortest paths lesson is the right tool.

#Grid Islands

A grid of land cells marked 1 and water cells marked 0 is a graph too. The items are cells numbered row * cols + col, and the links are land cells that touch. Every cell starts as its own group, so water cells are groups too. We subtract them at the end.

pythonrunnable cell

def count_islands(grid):
    """Count groups of connected '1' cells (up, down, left, right)."""
    rows, cols = len(grid), len(grid[0])
    uf = UnionFind(rows * cols)
    land = 0
    for r in range(rows):
        for c in range(cols):
            if grid[r][c] != "1":
                continue
            land += 1
            if r + 1 < rows and grid[r + 1][c] == "1":
                uf.union(r * cols + c, (r + 1) * cols + c)
            if c + 1 < cols and grid[r][c + 1] == "1":
                uf.union(r * cols + c, r * cols + c + 1)
    water = rows * cols - land
    return uf.components - water


grid = ["11000",
        "11000",
        "00100",
        "00011"]
print(count_islands(grid))
# Output:
# 3

We only look down and right, so each touching pair is joined once. The cost is O(R times C times α) time and O(R times C) space for R rows and C columns. A depth-first search does the same in O(R times C), but Union-Find wins when land cells are added one at a time and you want the island count after each.

#Accounts Merge

A site has several accounts, each with a name and some emails. Two accounts belong to one person if they share any email. The items are accounts. For each email we remember the first account that owned it, and every later account with that email is linked to that first one.

pythonrunnable cell

def merge_accounts(accounts):
    """accounts: [name, email, email, ...]. Merge accounts that share an email."""
    uf = UnionFind(len(accounts))
    owner = {}                                   # email -> first account that had it
    for i, (_name, *emails) in enumerate(accounts):
        for email in emails:
            if email in owner:
                uf.union(i, owner[email])
            else:
                owner[email] = i
    groups = {}
    for email, i in owner.items():
        groups.setdefault(uf.find(i), []).append(email)
    result = [[accounts[root][0]] + sorted(emails) for root, emails in groups.items()]
    return sorted(result)


accounts = [
    ["John", "john@a.com", "j@b.com"],
    ["Mary", "mary@c.com"],
    ["John", "j@b.com", "john2@d.com"],
    ["John", "other@e.com"],
]
for row in merge_accounts(accounts):
    print(row)
# Output:
# ['John', 'j@b.com', 'john2@d.com', 'john@a.com']
# ['John', 'other@e.com']
# ['Mary', 'mary@c.com']

Two of the three Johns shared j@b.com, so they merged. The third John has no shared email, so he stays alone, even though the name matches. With L emails in total, the unions cost O(L α) and the final sorting costs O(L log L), so the sort dominates. Space is O(L).

#Spanning Trees and the Cut Property

Now the second question: the cheapest network. A tree is a connected graph with no cycles. A spanning tree of a graph uses some of its edges to connect every node and still has no cycle. A connected graph with V nodes has spanning trees with exactly V-1 edges. Add one more edge and you make a cycle. Remove one and the tree falls into two groups.

A graph usually has many spanning trees. A minimum spanning tree, or MST, is a spanning tree whose total weight is smallest. Here is the graph we use from now on, with six towns A to F, numbered 0 to 5. Each triple is (weight, u, v).

text

A-B 1    A-C 3    B-C 4    B-D 2    C-D 5
C-E 7    D-E 6    D-F 8    E-F 9

Nine edges, so a spanning tree picks 5 of them. We can check every choice by brute force, using Union-Find to test each one. There are 126 ways to pick 5 edges from 9. Which of them are real spanning trees?

pythonrunnable cell

from itertools import combinations

EDGES = [(1, 0, 1), (3, 0, 2), (4, 1, 2), (2, 1, 3), (5, 2, 3),
         (7, 2, 4), (6, 3, 4), (8, 3, 5), (9, 4, 5)]       # (weight, u, v)


def is_spanning_tree(n, chosen):
    """True if the chosen edges join all n nodes with no cycle."""
    uf = UnionFind(n)
    return all(uf.union(u, v) for _w, u, v in chosen) and uf.components == 1


trees = [c for c in combinations(EDGES, 5) if is_spanning_tree(6, c)]
weights = sorted(sum(w for w, _u, _v in t) for t in trees)
print(len(list(combinations(EDGES, 5))), len(trees))
print(weights[:5], weights[-1])
print(weights.count(weights[0]))
# Output:
# 126 55
# [20, 21, 21, 21, 22] 31
# 1

Of 126 choices, 55 are spanning trees. Their weights run from 20 to 31, and exactly one tree has the smallest weight, 20. Brute force costs C(E, V-1) tests, which grows explosively, so we need a smarter way. It comes from one fact.

The cut property

A cut is any way to split the nodes into two non-empty sets. An edge crosses the cut if one end is on each side. The cut property says: look at the cheapest edge that crosses a cut. There is a minimum spanning tree that contains it. If that edge is strictly cheaper than every other crossing edge, then every MST contains it.

Why? The idea in plain words. Suppose some MST left that edge out. Add the edge to that tree. This creates a cycle, and a cycle that crosses the cut once must cross it again through some other edge. That other edge costs at least as much. Swap it for the cheap edge, and you get a tree that is no worse. So taking the cheapest crossing edge never hurts.

This is why greedy works here. A greedy algorithm makes the best-looking choice right now and never undoes it. Usually that is risky. For spanning trees the cut property proves it safe, and both algorithms below are just two ways to pick a cut.

#Kruskal's Algorithm

Kruskal looks at edges from cheapest to dearest. For each edge, it asks: are the two ends already in one group? If yes, the edge would close a cycle, so skip it. If no, take it and join the groups. This is exactly union returning True or False.

Why is a taken edge safe? At that moment, take the cut "the group of u against every other node". No edge taken so far crosses it, and every edge cheaper than this one has already been handled, so this is the cheapest edge crossing that cut. By the cut property, it belongs to an MST.

What Do You Think?

Kruskal has taken A-B (1), B-D (2) and A-C (3) on the town graph. The next edge in sorted order is B-C with weight 4. Does Kruskal take it or skip it?

pythonrunnable cell

NAMES = "ABCDEF"


def kruskal(n, edges):
    """Return (total weight, chosen edges). edges are (weight, u, v)."""
    uf = UnionFind(n)
    chosen, total = [], 0
    for w, u, v in sorted(edges):
        took = uf.union(u, v)
        print(f"  {w} {NAMES[u]}{NAMES[v]}", "take" if took else "skip (cycle)")
        if took:
            chosen.append((w, u, v))
            total += w
            if len(chosen) == n - 1:
                break                          # a spanning tree has n-1 edges
    return total, chosen


print(kruskal(6, EDGES))
# Output:
#   1 AB take
#   2 BD take
#   3 AC take
#   4 BC skip (cycle)
#   5 CD skip (cycle)
#   6 DE take
#   7 CE skip (cycle)
#   8 DF take
# (20, [(1, 0, 1), (2, 1, 3), (3, 0, 2), (6, 3, 4), (8, 3, 5)])

The result matches the brute-force minimum, 20, from five edges. The loop stopped after edge 8 and never looked at edge 9.

The cost. Sorting E edges is O(E log E) time. Since E is at most V times V, log E is at most 2 log V, so that is the same as O(E log V). The loop adds at most E union calls at O(α(V)) each, which is smaller. So Kruskal is O(E log V) time, dominated by the sort. Space is O(V) for Union-Find plus O(E) for the sorted copy of the edges.

#Prim's Algorithm

Prim grows one tree. Start at any node. Look at all edges from the tree to nodes outside it, take the cheapest, and add that new node. Repeat until every node is in. The cut is always "inside the tree" against "outside", so the cut property proves each choice safe.

To find the cheapest edge leaving the tree quickly, keep the candidate edges in a heap. The heap holds (weight, from, to) triples. When we pop a pair whose node is already inside, it is stale, so we skip it. This is the same trick you saw in Dijkstra's algorithm.

pythonrunnable cell

import heapq


def prim(n, edges, start=0):
    """Return (total weight, chosen edges). edges are (weight, u, v)."""
    adj = [[] for _ in range(n)]
    for w, u, v in edges:
        adj[u].append((w, v))
        adj[v].append((w, u))
    seen = [False] * n
    heap = [(0, start, start)]                 # (weight, from, to); the start has no real edge
    chosen, total = [], 0
    while heap:
        w, u, v = heapq.heappop(heap)
        if seen[v]:
            print(f"  pop {w} {NAMES[u]}{NAMES[v]} skip")
            continue
        seen[v] = True
        total += w
        if u != v:
            chosen.append((w, u, v))
            print(f"  pop {w} {NAMES[u]}{NAMES[v]} take")
        for w2, x in adj[v]:
            if not seen[x]:
                heapq.heappush(heap, (w2, v, x))
    return total, chosen


print(prim(6, EDGES))
# Output:
#   pop 1 AB take
#   pop 2 BD take
#   pop 3 AC take
#   pop 4 BC skip
#   pop 5 DC skip
#   pop 6 DE take
#   pop 7 CE skip
#   pop 8 DF take
#   pop 9 EF skip
# (20, [(1, 0, 1), (2, 1, 3), (3, 0, 2), (6, 3, 4), (8, 3, 5)])

Same tree and same weight as Kruskal. Notice the order of the taken edges is the same here, but that is a coincidence of this graph. Prim always grows from the start, while Kruskal may build several pieces that join later.

The cost. Each edge is pushed onto the heap at most once, when its first end enters the tree, so the heap never holds more than E entries. Each push and pop costs O(log E), which is O(log V). Total time is O(E log V). Space is O(V + E): the adjacency list is O(V + E), the heap at most O(E), and seen is O(V).

If the graph is not connected, Prim only reaches the start's group, and Kruskal ends with a spanning forest of fewer than V-1 edges. Check len(chosen) == n - 1 before you trust the answer.

#Kruskal or Prim?

Both give a tree of the same minimum weight, and with these implementations both cost O(E log V). So the choice is about the shape of your data and what you need.

Question	Kruskal	Prim
Natural input	an edge list	an adjacency list
Works from	all edges, sorted	one growing tree
Needs	`sort` and Union-Find	`heapq` and `seen`
Disconnected graph	gives a minimum spanning forest	reaches only the start's group
Early stop	can stop after V-1 edges are taken	stops when all nodes are in

Sparse versus dense. A sparse graph has few edges, E close to V. A dense graph has many, E close to V times V. On a sparse graph both are about O(V log V), and Kruskal's sort of a short list is simple and fast.

On a dense graph, E log V is about V² log V. A Prim version that drops the heap, and instead scans a plain array of "cheapest known edge to each outside node" in every round, costs O(V²), and that is better when E is close to V². Running that version on the town graph also gives 20.

The next cell counts the work on a sparse and a dense graph with 40 nodes. Kruskal's sort always touches all E edges, but the loop can stop early. Prim's heap gets each edge once.

pythonrunnable cell

import heapq


def kruskal_work(n, edges):
    """Return (total weight, edges looked at after sorting)."""
    uf = UnionFind(n)
    total = looked = taken = 0
    for w, u, v in sorted(edges):
        looked += 1
        if uf.union(u, v):
            total += w
            taken += 1
            if taken == n - 1:
                break
    return total, looked


def prim_work(n, edges):
    """Return (total weight, heap pushes)."""
    adj = [[] for _ in range(n)]
    for w, u, v in edges:
        adj[u].append((w, v))
        adj[v].append((w, u))
    seen = [False] * n
    heap = [(0, 0)]
    total = pushes = 0
    while heap:
        w, v = heapq.heappop(heap)
        if seen[v]:
            continue
        seen[v] = True
        total += w
        for w2, x in adj[v]:
            if not seen[x]:
                heapq.heappush(heap, (w2, x))
                pushes += 1
    return total, pushes


n = 40
dense = [((i * 31 + j * 17) % 97 + 1, i, j) for i in range(n) for j in range(i + 1, n)]
sparse = [((i * 31 + 5) % 97 + 1, i, i + 1) for i in range(n - 1)]
sparse += [((i * 13) % 97 + 1, i, (i + 7) % n) for i in range(n) if i < (i + 7) % n]
for label, edges in (("sparse", sparse), ("dense", dense)):
    print(label, len(edges), kruskal_work(n, edges), prim_work(n, edges))
# Output:
# sparse 72 (1001, 50) (1001, 72)
# dense 780 (115, 47) (115, 780)

Read each line as: edge count, then (weight, work) for Kruskal, then (weight, work) for Prim. The weights agree. On the dense graph Kruskal looked at 47 of 780 sorted edges, but it still paid to sort all 780. Prim pushed all 780 edges.

#Is the Minimum Spanning Tree Unique?

The weight of a minimum spanning tree is unique. It is just the smallest total, one number, so there cannot be two different smallest totals. The tree itself may not be. Three towns with all three cables priced 1 have three different trees of weight 2. Kruskal picks among ties by the order of the sorted list, so changing the order changes the tree but never the weight.

pythonrunnable cell

def kruskal_edges(n, edges):
    """Kruskal on edges (weight, u, v). Return the chosen (u, v) pairs. Ties keep input order."""
    uf = UnionFind(n)
    return [(u, v) for _w, u, v in sorted(edges, key=lambda e: e[0]) if uf.union(u, v)]


triangle = [(1, 0, 1), (1, 1, 2), (1, 0, 2)]
print(kruskal_edges(3, triangle))
print(kruskal_edges(3, triangle[::-1]))
# Output:
# [(0, 1), (1, 2)]
# [(0, 2), (1, 2)]

Both trees weigh 2.

Now the other direction. If all edge weights are different, there is exactly one MST. The reason, in short. Suppose two MSTs differ, and look at the cheapest edge that is in one but not the other, say in the first. Add it to the second tree. That makes a cycle, and the cycle has some edge that is not in the first tree. That edge is strictly dearer, because weights differ. Swapping it out gives a tree cheaper than the second, which was supposed to be minimum. That is a contradiction, so there are not two.

On the town graph all nine weights differ, and the brute-force count above found exactly one tree of weight 20.

#In Machine Learning: Single-Linkage Clustering

Clustering groups points that are near each other. Single-linkage clustering starts with every point as its own cluster. It repeatedly merges the two clusters that have the closest pair of points, one from each, until the number of clusters you asked for remains. Look at it as a graph: the points are nodes, and the edge between any two points weighs their distance.

Merging the closest pair of clusters first is exactly Kruskal's order. Sort the distances, union the two points if they are in different clusters, and stop when components equals k. The tree Kruskal builds along the way is the MST of the points, and the distances at which merges happen are its edge weights. Libraries such as SciPy offer single linkage as one of their hierarchical clustering methods.

pythonrunnable cell

import math

points = [(1.0, 1.0), (1.5, 2.0), (2.0, 1.2), (8.0, 8.0),
          (8.5, 7.5), (9.0, 8.4), (15.0, 1.0), (15.5, 1.8)]
n = len(points)
edges = sorted((math.dist(points[i], points[j]), i, j)
               for i in range(n) for j in range(i + 1, n))


def single_linkage(k):
    """Merge the two closest clusters until k remain. Return the clusters and merge distances."""
    uf = UnionFind(n)
    merges = []
    for d, i, j in edges:
        if uf.components == k:
            break
        if uf.union(i, j):
            merges.append(round(d, 2))
    groups = {}
    for i in range(n):
        groups.setdefault(uf.find(i), []).append(i)
    return sorted(groups.values()), merges


for k in (3, 2):
    print(k, single_linkage(k)[0])
print(single_linkage(1)[1])
# Output:
# 3 [[0, 1, 2], [3, 4, 5], [6, 7]]
# 2 [[0, 1, 2, 3, 4, 5], [6, 7]]
# [0.71, 0.94, 0.94, 1.02, 1.03, 8.85, 9.03]

The last line is all seven MST edges. Five are tiny, and then come two big jumps, 8.85 and 9.03. Asking for 3 clusters is the same as cutting the two heaviest MST edges. A big jump between merge distances is a common hint about how many clusters the data has.

The cost for n points is high on purpose: there are n(n-1)/2 distances, so sorting them is O(n² log n) time and O(n²) space. Single linkage also has a known weakness. A thin chain of points between two blobs can glue them into one cluster, because only the closest pair matters.

#When Not to Use Union-Find

Union-Find only joins. It cannot split a group apart or delete a link, because compression has thrown the old shape away. If links can be removed, use another tool, or process the changes backwards in time so every removal becomes an addition.

It also gives a yes or no, not a path. It tells you two towns are connected and not how to get from one to the other. Use BFS or DFS for a path, and Dijkstra when edges have weights and you want the cheapest route between two nodes. A spanning tree is not that: it is the cheapest way to join everything, and the path inside it between two towns is usually not the shortest one.

For a graph that never changes and you only need the components once, one DFS in O(V + E) is just as good, and simpler. Finally, the MST in this lesson is for undirected graphs. A directed graph needs a different algorithm.

#Common Mistakes

Watch Out

Linking the items instead of their roots

Fix: Always call find on both items first, and hang one root under the other.

pythonrunnable cell

parent = list(range(4))


def find(x):
    while parent[x] != x:
        x = parent[x]
    return x


def bad_union(a, b):
    parent[a] = b          # a may not be a root


bad_union(0, 1)
bad_union(0, 2)
print(find(1) == find(2))
# Output:
# False

The first call linked 0 to 1. The second call overwrote that link, so 0 now points at 2, and item 1 was dropped from the group. Items 0, 1 and 2 were all supposed to be together. Joining the roots keeps every earlier link alive.

Watch Out

Trusting a spanning tree on a graph that is not connected

Fix: After Kruskal or Prim, check that you took exactly n - 1 edges. Fewer means some nodes could not be reached.

pythonrunnable cell

n = 4
edges = [(1, 0, 1), (2, 2, 3)]          # two separate pairs
uf = UnionFind(n)
chosen = [e for e in sorted(edges) if uf.union(e[1], e[2])]
print(sum(w for w, _u, _v in chosen), len(chosen), len(chosen) == n - 1)
# Output:
# 3 2 False

The code returned a total of 3 as if it were a network, but it joined only two pairs of towns. A spanning tree of four nodes needs 3 edges, and we have 2. Always test the edge count, or uf.components == 1, before you report a cost.

Quick check

In Prim's code, the first line after `heappop` is `if seen[v]: continue`. What would go wrong without it?

#Try It: The Five-Problem Ladder

Five problems, easy to hard. The class from the lesson is at the top of the cell. Problem 1 changes one thing about the island example. Problem 2 is a small variation on cycle detection. Problem 3 combines Union-Find with sorting to build a network. Problem 4 has two bugs to find. Problem 5 is slow, and you make it fast and say its Big-O. Predict each printed line before you run.

pythonplayground.py · Pyodide

# Five problems, easy to hard. Fill in each TODO, then run.
import heapq

class UnionFind:
  """The class from the lesson, ready to use."""

def __init__(self, n):
      self.parent = list(range(n))
      self.rank = [0] * n
      self.components = n

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

def union(self, x, y):
      root_x, root_y = self.find(x), self.find(y)
      if root_x == root_y:
          return False
      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
      self.components -= 1
      return True

def connected(self, x, y):
      return self.find(x) == self.find(y)

# 1. Islands, with diagonals. A land cell touches its 4 side neighbours, and when
#    diagonals is True it also touches its 4 corner neighbours. Count the islands.
def count_islands(grid, diagonals):
  rows, cols = len(grid), len(grid[0])
  uf = UnionFind(rows * cols)
  land = 0
  # TODO: build steps, the list of (row change, column change) pairs to look at:
  #       4 for the sides, and 4 more for the corners when diagonals is True
  steps = []
  for r in range(rows):
      for c in range(cols):
          if grid[r][c] != "1":
              continue
          land += 1
          # TODO: for each step, if the neighbour is inside the grid and is land, union the two cells
  return uf.components - (rows * cols - land)

# 2. The extra link. The edges below would form a tree, except that one extra edge was added.
#    Return the first edge, in the order given, that joins two nodes already connected.
def redundant_edge(n, edges):
  uf = UnionFind(n)
  # TODO: union each edge; return the edge when union says False; return None if none does
  return None

# 3. Cheapest network. roads is a list of (cost, a, b). Return the least total cost that
#    connects all n cities, or -1 if some city cannot be reached at all.
def cheapest_network(n, roads):
  uf = UnionFind(n)
  total = 0
  # TODO: Kruskal. Go through the roads cheapest first and take a road only if union returns True.
  #       At the end, answer -1 unless exactly one group is left.
  return total

# 4. Find the bugs. The groups should merge, and components should count groups.
#    There are two bugs, both in union.
class Groups:
  def __init__(self, n):
      self.parent = list(range(n))
      self.components = n

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

def union(self, x, y):
      self.parent[x] = y
      self.components -= 1
      return True

def connected(self, x, y):
      return self.find(x) == self.find(y)

# 5. Make it faster. This finds the weight of a minimum spanning tree by scanning every edge
#    in every round, which is O(V * E). Write mst_weight_fast with Prim and heapq, which is O(E log V).
def mst_weight_slow(n, edges):
  inside = {0}
  total = 0
  while len(inside) < n:
      best = None
      for w, u, v in edges:
          if (u in inside) != (v in inside) and (best is None or w < best[0]):
              best = (w, u, v)
      if best is None:
          return -1
      total += best[0]
      inside.update((best[1], best[2]))
  return total

def mst_weight_fast(n, edges):
  adj = [[] for _ in range(n)]
  # TODO: fill adj, then run Prim with a heap of (cost, vertex); skip a vertex already seen;
  #       return the total, or -1 if fewer than n vertices were reached
  return None

grid = ["10100",
      "01000",
      "00011"]
print(count_islands(grid, False), count_islands(grid, True))

print(redundant_edge(5, [(0, 1), (1, 2), (2, 3), (3, 1), (0, 4)]))
print(redundant_edge(3, [(0, 1), (1, 2)]))

roads = [(5, 0, 1), (3, 1, 2), (4, 0, 2), (1, 2, 3)]
print(cheapest_network(4, roads), cheapest_network(5, roads))

g = Groups(5)
g.union(0, 1)
g.union(0, 2)
g.union(0, 2)
print(g.components, g.connected(1, 2), g.connected(0, 4))

edges = [(1, 0, 1), (3, 0, 2), (4, 1, 2), (2, 1, 3), (5, 2, 3), (7, 2, 4), (6, 3, 4), (8, 3, 5), (9, 4, 5)]
print(mst_weight_slow(6, edges), mst_weight_fast(6, edges))
big = [((i * 31 + j * 17) % 97 + 1, i, j) for i in range(40) for j in range(i + 1, 40)]
print(mst_weight_slow(40, big), mst_weight_fast(40, big))

Tests · Correct output, line by line: '4 2'; '(3, 1)'; 'None'; '8 -1'; '3 True False'; '20 20'; '115 115'. With the TODOs left as they are, you see '5 5', None, None, '0 0', '2 False False', '20 None', '115 None'.

Costs for the ladder, with reasons. Problem 1 is O(R times C times α) time, because each cell does at most 8 unions, and O(R times C) space for the parent list.

Problem 2 is O(n + E α) time and O(n) space. Problem 3 is O(E log E) time for the sort and O(n + E) space. Problem 4 is a correctness fix only. Problem 5 drops from O(V times E), because each of V rounds scans E edges, to O(E log V) time and O(V + E) space.

#Recap

Union-Find keeps items in groups using a parent list. find climbs to a root, and union hangs one root under another. Path compression flattens every path that find climbs. Union by rank hangs the shorter tree under the taller one, so a tree is never more than log n levels tall. Together they make each operation O(α(n)) amortized, where the inverse Ackermann function α stays below 5 for any real input, so it acts like a constant.

It answers dynamic components, cycles in undirected graphs, grid islands and merged accounts. A spanning tree joins every node with V-1 edges and no cycle. The cut property says the cheapest edge across any cut is safe to take. Kruskal sorts edges and uses Union-Find, in O(E log V) time. Prim grows a tree with a heap, also in O(E log V). The MST weight is unique, and single-linkage clustering is Kruskal with a stop.

Quick Check1 / 6

With the plain parent array and no tricks, what does this print? parent = [0, 0, 2, 2, 4, 5] root_a = 3 while parent[root_a] != root_a: root_a = parent[root_a] print(root_a)

Recap

Key Takeaways

Union-Find stores groups in a parent list: find climbs to a root, and union hangs one root under another
Path compression flattens each path find climbs, and union by rank limits height to log n; together they give O(α(n)) amortized per operation
Use it for streams of new links, cycles in undirected graphs, grid islands and merged accounts, and remember that it cannot split a group
A spanning tree of V nodes has V-1 edges, and the cut property says the cheapest edge across any cut is safe to take
Kruskal sorts the edges and keeps an edge when union returns True, in O(E log V) time and O(V + E) space
Prim grows one tree with a heap and skips stale entries, in O(E log V) time and O(V + E) space; the MST weight is unique even when the tree is not

Union-Find and Minimum Spanning Trees

Learning Objectives

After this lesson, you will be able to:

Explain the parent array and trace find and union by hand on a small example
Say what path compression and union by rank each do, and show with counted steps why the two together are almost O(1) per operation
State in plain words what the inverse Ackermann function is, and where the formal amortized argument lives
Use Union-Find for a stream of new links, for a cycle in an undirected graph, for grid islands and for merging accounts
Say what a spanning tree is, state the cut property in plain words, and say why it makes greedy choices safe
Write Kruskal with sort and Union-Find, write Prim with heapq, give the time and space of each, and choose between them
Explain why the weight of a minimum spanning tree is unique, when the tree itself is unique, and how single-linkage clustering is Kruskal in disguise

Prerequisites

Before You Start

Review

#Two Questions About Connections

#Quick Recap: Graphs and Heaps

Five lines of what you need.

A graph has V nodes and E edges. An edge can carry a weight.
An edge list is a plain list of (weight, u, v) triples. An adjacency list gives each node its neighbours (Graphs).
A cycle is a path that leaves a node and returns to it without reusing an edge.
A heap hands you the smallest value in O(log n), and heapq is Python's version (Heaps).
sorted costs O(n log n) time and O(n) extra space (Sorting Algorithms).

#The Parent Array Idea

To find the group of an item, follow the parent links upward until you reach a root. That operation is called find. Two items are in the same group exactly when their roots are the same item.

To join two groups, find both roots and hang one root under the other. That operation is called union. We only ever change the parent of a root. Trace it by hand first.

Step	Call	parent after the call	Groups
start		`[0, 1, 2, 3, 4, 5]`	six groups of one
1	`union(0, 1)`	`[0, 0, 2, 3, 4, 5]`	1 goes under 0
2	`union(2, 3)`	`[0, 0, 2, 2, 4, 5]`	3 goes under 2
3	`union(1, 3)`	`[0, 0, 0, 2, 4, 5]`	root of 1 is 0, root of 3 is 2, so 2 goes under 0

text

   0       4     5
  / \
 1   2
     |
     3

Now the code. It has the three pieces from the table: a list, a find that climbs, and a union that links two roots.

pythonrunnable cell

parent = list(range(6))          # every item starts as its own root


def find(x):
    """Walk up the parent links until we reach a root."""
    while parent[x] != x:
        x = parent[x]
    return x


def union(a, b):
    """Join the groups of a and b by hanging one root under the other."""
    root_a, root_b = find(a), find(b)
    if root_a != root_b:
        parent[root_b] = root_a


union(0, 1)
print(parent)
union(2, 3)
print(parent)
union(1, 3)
print(parent)
print(find(3), find(4))
print(find(1) == find(3), find(1) == find(4))
# Output:
# [0, 0, 2, 3, 4, 5]
# [0, 0, 2, 2, 4, 5]
# [0, 0, 0, 2, 4, 5]
# 0 4
# True False

The printed lists match the table row by row. Item 3 climbs 3, then 2, then 0, so its root is 0. Item 4 never moved, so it is its own root.

What Do You Think?

Start from parent = [0, 0, 2, 2, 4, 5] (the state after step 2 above). The code is the plain version: it hangs the second root under the first. What is parent after union(3, 5)?

#Why the Plain Version Gets Slow

pythonrunnable cell

def chain_hops(n):
    """Plain union only. Join item i+1 to the group of item i, for every i."""
    parent = list(range(n))

    def find(x):
        hops = 0
        while parent[x] != x:
            x = parent[x]
            hops += 1
        return x, hops

    for i in range(n - 1):
        root_a, _ = find(i + 1)
        root_b, _ = find(i)
        if root_a != root_b:
            parent[root_b] = root_a
    return find(0)[1]


for n in (10, 100, 1_000):
    print(n, chain_hops(n))
# Output:
# 10 9
# 100 99
# 1000 999

#Trick One: Path Compression

pythonrunnable cell

parent = [1, 2, 3, 4, 4]          # a chain: 0 -> 1 -> 2 -> 3 -> 4, and 4 is the root


def find(x):
    """Find the root, then point every item on the path straight at it."""
    root = x
    while parent[root] != root:       # pass 1: walk up to the root
        root = parent[root]
    while parent[x] != root:          # pass 2: re-point each item on the path
        parent[x], x = root, parent[x]
    return root


print(parent)
print(find(0))
print(parent)
# Output:
# [1, 2, 3, 4, 4]
# 4
# [4, 4, 4, 4, 4]

#Trick Two: Union by Rank

What Do You Think?

pythonrunnable cell

def deepest(n, pairs):
    """Union by rank, no compression. Return how many steps the deepest item is from its root."""
    parent = list(range(n))
    rank = [0] * n

    def find(x):
        while parent[x] != x:
            x = parent[x]
        return x

    for a, b in pairs:
        ra, rb = find(a), find(b)
        if ra == rb:
            continue
        if rank[ra] < rank[rb]:
            ra, rb = rb, ra
        if rank[ra] == rank[rb]:
            rank[ra] += 1
        parent[rb] = ra
    depth = 0
    for x in range(n):
        steps, y = 0, x
        while parent[y] != y:
            y, steps = parent[y], steps + 1
        depth = max(depth, steps)
    return depth


def pairs_of_equal_trees(n):
    """Join items in pairs, then join the pairs, then join those, and so on."""
    pairs, step = [], 1
    while step < n:
        pairs += [(i, i + step) for i in range(0, n, 2 * step)]
        step *= 2
    return pairs


for n in (16, 256, 4096):
    print(n, deepest(n, pairs_of_equal_trees(n)))
# Output:
# 16 4
# 256 8
# 4096 12

#The Finished UnionFind Class

pythonrunnable cell

class UnionFind:
    """Groups of items numbered 0 to n-1, with path compression and union by rank."""

    def __init__(self, n: int) -> None:
        """Start with n groups of one item each."""
        self.parent: list[int] = list(range(n))   # each item is its own root
        self.rank: list[int] = [0] * n            # upper bound on tree height
        self.components: int = n                  # how many groups there are now

    def find(self, x: int) -> int:
        """Return the root of x's group, flattening the path on the way."""
        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, x: int, y: int) -> bool:
        """Merge the groups of x and y. Return True if they were different groups."""
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        if self.rank[root_x] < self.rank[root_y]:  # union by rank:
            root_x, root_y = root_y, root_x        # the taller tree stays on top
        self.parent[root_y] = root_x
        if self.rank[root_x] == self.rank[root_y]:
            self.rank[root_x] += 1
        self.components -= 1
        return True

    def connected(self, x: int, y: int) -> bool:
        """Say whether x and y are in the same group."""
        return self.find(x) == self.find(y)


uf = UnionFind(5)
print(uf.parent, uf.components)
uf.union(0, 1)
print(uf.parent, uf.components)
uf.union(2, 3)
print(uf.parent, uf.components)
uf.union(1, 3)
print(uf.parent, uf.rank, uf.components)
print(uf.find(3), uf.parent)
print(uf.connected(1, 3), uf.connected(1, 4))
print(uf.union(0, 3))
# Output:
# [0, 1, 2, 3, 4] 5
# [0, 0, 2, 3, 4] 4
# [0, 0, 2, 2, 4] 3
# [0, 0, 0, 2, 4] [2, 0, 1, 0, 0] 2
# 0 [0, 0, 0, 0, 4]
# True False
# False

#How Fast Is It Really?

pythonrunnable cell

def total_hops(n, pairs, compress, by_rank):
    """Do all the unions, then one find per item. Count every step up a parent link."""
    parent = list(range(n))
    rank = [0] * n
    hops = 0

    def find(x):
        nonlocal hops
        root = x
        while parent[root] != root:
            root = parent[root]
            hops += 1
        if compress:
            while parent[x] != root:
                parent[x], x = root, parent[x]
        return root

    for a, b in pairs:
        ra, rb = find(a), find(b)
        if ra == rb:
            continue
        if by_rank:
            if rank[ra] < rank[rb]:
                ra, rb = rb, ra
            if rank[ra] == rank[rb]:
                rank[ra] += 1
        parent[rb] = ra
    for x in range(n):
        find(x)
    return hops


n = 4096
chain = [(i + 1, i) for i in range(n - 1)]
mixed = [(i, (i * 7919 + 13) % n) for i in range(n)]
print("tricks            chain   mixed")
for label, compress, by_rank in (("neither", False, False),
                                 ("compression only", True, False),
                                 ("rank only", False, True),
                                 ("both", True, True)):
    print(f"{label:<16} {total_hops(n, chain, compress, by_rank):>7} {total_hops(n, mixed, compress, by_rank):>7}")

print()
print("both tricks, mixed workload")
for n in (1_000, 10_000, 100_000):
    pairs = [(i, (i * 7919 + 13) % n) for i in range(n)]
    hops = total_hops(n, pairs, True, True)
    print(n, hops, f"{hops / (3 * n):.2f} hops per find")
# Output:
# tricks            chain   mixed
# neither          8386560   26553
# compression only    8189   12019
# rank only           8188   18565
# both                8188   11983
#
# both tricks, mixed workload
# 1000 2426 0.81 hops per find
# 10000 27721 0.92 hops per find
# 100000 295450 0.98 hops per find

Now the plain-words version of the famous bound.

Quick check

You keep path compression but remove union by rank, so the root of the second group always goes under the root of the first. Which statement is true about the cost of an operation?

#Using Union-Find: Four Classic Jobs

All four jobs below use the same class. Each time, ask: what are the items, and what counts as a link?

#Groups That Change Over Time

pythonrunnable cell

uf = UnionFind(6)
for a, b in [(0, 1), (2, 3), (1, 3), (4, 5), (0, 2)]:
    merged = uf.union(a, b)
    print((a, b), "merged" if merged else "already together", uf.components)
# Output:
# (0, 1) merged 5
# (2, 3) merged 4
# (1, 3) merged 3
# (4, 5) merged 2
# (0, 2) already together 2

The last link joined two people who were already in one group, so the count stayed at 2. For m links over n people, the total is O((n + m) α(n)) time and O(n) space.

#A Cycle in an Undirected Graph

pythonrunnable cell

def has_cycle(n, edges):
    """True if the undirected graph has a cycle. O(n + E * alpha(n)) time, O(n) space."""
    uf = UnionFind(n)
    for u, v in edges:
        if not uf.union(u, v):       # both ends already joined, so this edge closes a loop
            return True
    return False


print(has_cycle(4, [(0, 1), (1, 2), (2, 3)]))
print(has_cycle(4, [(0, 1), (1, 2), (2, 3), (3, 1)]))
print(has_cycle(3, [(0, 1), (1, 1)]))
# Output:
# False
# True
# True

#Grid Islands

pythonrunnable cell

def count_islands(grid):
    """Count groups of connected '1' cells (up, down, left, right)."""
    rows, cols = len(grid), len(grid[0])
    uf = UnionFind(rows * cols)
    land = 0
    for r in range(rows):
        for c in range(cols):
            if grid[r][c] != "1":
                continue
            land += 1
            if r + 1 < rows and grid[r + 1][c] == "1":
                uf.union(r * cols + c, (r + 1) * cols + c)
            if c + 1 < cols and grid[r][c + 1] == "1":
                uf.union(r * cols + c, r * cols + c + 1)
    water = rows * cols - land
    return uf.components - water


grid = ["11000",
        "11000",
        "00100",
        "00011"]
print(count_islands(grid))
# Output:
# 3

#Accounts Merge

pythonrunnable cell

def merge_accounts(accounts):
    """accounts: [name, email, email, ...]. Merge accounts that share an email."""
    uf = UnionFind(len(accounts))
    owner = {}                                   # email -> first account that had it
    for i, (_name, *emails) in enumerate(accounts):
        for email in emails:
            if email in owner:
                uf.union(i, owner[email])
            else:
                owner[email] = i
    groups = {}
    for email, i in owner.items():
        groups.setdefault(uf.find(i), []).append(email)
    result = [[accounts[root][0]] + sorted(emails) for root, emails in groups.items()]
    return sorted(result)


accounts = [
    ["John", "john@a.com", "j@b.com"],
    ["Mary", "mary@c.com"],
    ["John", "j@b.com", "john2@d.com"],
    ["John", "other@e.com"],
]
for row in merge_accounts(accounts):
    print(row)
# Output:
# ['John', 'j@b.com', 'john2@d.com', 'john@a.com']
# ['John', 'other@e.com']
# ['Mary', 'mary@c.com']

#Spanning Trees and the Cut Property

text

A-B 1    A-C 3    B-C 4    B-D 2    C-D 5
C-E 7    D-E 6    D-F 8    E-F 9

pythonrunnable cell

from itertools import combinations

EDGES = [(1, 0, 1), (3, 0, 2), (4, 1, 2), (2, 1, 3), (5, 2, 3),
         (7, 2, 4), (6, 3, 4), (8, 3, 5), (9, 4, 5)]       # (weight, u, v)


def is_spanning_tree(n, chosen):
    """True if the chosen edges join all n nodes with no cycle."""
    uf = UnionFind(n)
    return all(uf.union(u, v) for _w, u, v in chosen) and uf.components == 1


trees = [c for c in combinations(EDGES, 5) if is_spanning_tree(6, c)]
weights = sorted(sum(w for w, _u, _v in t) for t in trees)
print(len(list(combinations(EDGES, 5))), len(trees))
print(weights[:5], weights[-1])
print(weights.count(weights[0]))
# Output:
# 126 55
# [20, 21, 21, 21, 22] 31
# 1

The cut property

#Kruskal's Algorithm

What Do You Think?

Kruskal has taken A-B (1), B-D (2) and A-C (3) on the town graph. The next edge in sorted order is B-C with weight 4. Does Kruskal take it or skip it?

pythonrunnable cell

NAMES = "ABCDEF"


def kruskal(n, edges):
    """Return (total weight, chosen edges). edges are (weight, u, v)."""
    uf = UnionFind(n)
    chosen, total = [], 0
    for w, u, v in sorted(edges):
        took = uf.union(u, v)
        print(f"  {w} {NAMES[u]}{NAMES[v]}", "take" if took else "skip (cycle)")
        if took:
            chosen.append((w, u, v))
            total += w
            if len(chosen) == n - 1:
                break                          # a spanning tree has n-1 edges
    return total, chosen


print(kruskal(6, EDGES))
# Output:
#   1 AB take
#   2 BD take
#   3 AC take
#   4 BC skip (cycle)
#   5 CD skip (cycle)
#   6 DE take
#   7 CE skip (cycle)
#   8 DF take
# (20, [(1, 0, 1), (2, 1, 3), (3, 0, 2), (6, 3, 4), (8, 3, 5)])

The result matches the brute-force minimum, 20, from five edges. The loop stopped after edge 8 and never looked at edge 9.

#Prim's Algorithm

pythonrunnable cell

import heapq


def prim(n, edges, start=0):
    """Return (total weight, chosen edges). edges are (weight, u, v)."""
    adj = [[] for _ in range(n)]
    for w, u, v in edges:
        adj[u].append((w, v))
        adj[v].append((w, u))
    seen = [False] * n
    heap = [(0, start, start)]                 # (weight, from, to); the start has no real edge
    chosen, total = [], 0
    while heap:
        w, u, v = heapq.heappop(heap)
        if seen[v]:
            print(f"  pop {w} {NAMES[u]}{NAMES[v]} skip")
            continue
        seen[v] = True
        total += w
        if u != v:
            chosen.append((w, u, v))
            print(f"  pop {w} {NAMES[u]}{NAMES[v]} take")
        for w2, x in adj[v]:
            if not seen[x]:
                heapq.heappush(heap, (w2, v, x))
    return total, chosen


print(prim(6, EDGES))
# Output:
#   pop 1 AB take
#   pop 2 BD take
#   pop 3 AC take
#   pop 4 BC skip
#   pop 5 DC skip
#   pop 6 DE take
#   pop 7 CE skip
#   pop 8 DF take
#   pop 9 EF skip
# (20, [(1, 0, 1), (2, 1, 3), (3, 0, 2), (6, 3, 4), (8, 3, 5)])

If the graph is not connected, Prim only reaches the start's group, and Kruskal ends with a spanning forest of fewer than V-1 edges. Check len(chosen) == n - 1 before you trust the answer.

#Kruskal or Prim?

Both give a tree of the same minimum weight, and with these implementations both cost O(E log V). So the choice is about the shape of your data and what you need.

Question	Kruskal	Prim
Natural input	an edge list	an adjacency list
Works from	all edges, sorted	one growing tree
Needs	`sort` and Union-Find	`heapq` and `seen`
Disconnected graph	gives a minimum spanning forest	reaches only the start's group
Early stop	can stop after V-1 edges are taken	stops when all nodes are in

The next cell counts the work on a sparse and a dense graph with 40 nodes. Kruskal's sort always touches all E edges, but the loop can stop early. Prim's heap gets each edge once.

pythonrunnable cell

import heapq


def kruskal_work(n, edges):
    """Return (total weight, edges looked at after sorting)."""
    uf = UnionFind(n)
    total = looked = taken = 0
    for w, u, v in sorted(edges):
        looked += 1
        if uf.union(u, v):
            total += w
            taken += 1
            if taken == n - 1:
                break
    return total, looked


def prim_work(n, edges):
    """Return (total weight, heap pushes)."""
    adj = [[] for _ in range(n)]
    for w, u, v in edges:
        adj[u].append((w, v))
        adj[v].append((w, u))
    seen = [False] * n
    heap = [(0, 0)]
    total = pushes = 0
    while heap:
        w, v = heapq.heappop(heap)
        if seen[v]:
            continue
        seen[v] = True
        total += w
        for w2, x in adj[v]:
            if not seen[x]:
                heapq.heappush(heap, (w2, x))
                pushes += 1
    return total, pushes


n = 40
dense = [((i * 31 + j * 17) % 97 + 1, i, j) for i in range(n) for j in range(i + 1, n)]
sparse = [((i * 31 + 5) % 97 + 1, i, i + 1) for i in range(n - 1)]
sparse += [((i * 13) % 97 + 1, i, (i + 7) % n) for i in range(n) if i < (i + 7) % n]
for label, edges in (("sparse", sparse), ("dense", dense)):
    print(label, len(edges), kruskal_work(n, edges), prim_work(n, edges))
# Output:
# sparse 72 (1001, 50) (1001, 72)
# dense 780 (115, 47) (115, 780)

#Is the Minimum Spanning Tree Unique?

pythonrunnable cell

def kruskal_edges(n, edges):
    """Kruskal on edges (weight, u, v). Return the chosen (u, v) pairs. Ties keep input order."""
    uf = UnionFind(n)
    return [(u, v) for _w, u, v in sorted(edges, key=lambda e: e[0]) if uf.union(u, v)]


triangle = [(1, 0, 1), (1, 1, 2), (1, 0, 2)]
print(kruskal_edges(3, triangle))
print(kruskal_edges(3, triangle[::-1]))
# Output:
# [(0, 1), (1, 2)]
# [(0, 2), (1, 2)]

Both trees weigh 2.

On the town graph all nine weights differ, and the brute-force count above found exactly one tree of weight 20.

#In Machine Learning: Single-Linkage Clustering

pythonrunnable cell

import math

points = [(1.0, 1.0), (1.5, 2.0), (2.0, 1.2), (8.0, 8.0),
          (8.5, 7.5), (9.0, 8.4), (15.0, 1.0), (15.5, 1.8)]
n = len(points)
edges = sorted((math.dist(points[i], points[j]), i, j)
               for i in range(n) for j in range(i + 1, n))


def single_linkage(k):
    """Merge the two closest clusters until k remain. Return the clusters and merge distances."""
    uf = UnionFind(n)
    merges = []
    for d, i, j in edges:
        if uf.components == k:
            break
        if uf.union(i, j):
            merges.append(round(d, 2))
    groups = {}
    for i in range(n):
        groups.setdefault(uf.find(i), []).append(i)
    return sorted(groups.values()), merges


for k in (3, 2):
    print(k, single_linkage(k)[0])
print(single_linkage(1)[1])
# Output:
# 3 [[0, 1, 2], [3, 4, 5], [6, 7]]
# 2 [[0, 1, 2, 3, 4, 5], [6, 7]]
# [0.71, 0.94, 0.94, 1.02, 1.03, 8.85, 9.03]

#When Not to Use Union-Find

#Common Mistakes

Watch Out

Linking the items instead of their roots

Fix: Always call find on both items first, and hang one root under the other.

pythonrunnable cell

parent = list(range(4))


def find(x):
    while parent[x] != x:
        x = parent[x]
    return x


def bad_union(a, b):
    parent[a] = b          # a may not be a root


bad_union(0, 1)
bad_union(0, 2)
print(find(1) == find(2))
# Output:
# False

Watch Out

Trusting a spanning tree on a graph that is not connected

Fix: After Kruskal or Prim, check that you took exactly n - 1 edges. Fewer means some nodes could not be reached.

pythonrunnable cell

n = 4
edges = [(1, 0, 1), (2, 2, 3)]          # two separate pairs
uf = UnionFind(n)
chosen = [e for e in sorted(edges) if uf.union(e[1], e[2])]
print(sum(w for w, _u, _v in chosen), len(chosen), len(chosen) == n - 1)
# Output:
# 3 2 False