Practice Problems

Each of the four problems in this chapter captures one classic use of union-find. Every problem comes with an approach and a Python solution. The solutions share the short DSU below.

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    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
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True

Problem 1. Counting clubs

You are given n students (numbered 0 to n-1) and a list of pairs meaning "these two students are in the same club". Club membership is transitive. Find how many clubs there are and how many members the largest club has.

Approach: call union for every pair, then count the distinct roots and take the maximum size stored at those roots.

def clubs(n, pairs):
    dsu = DSU(n)
    for a, b in pairs:
        dsu.union(a, b)
    roots = {dsu.find(x) for x in range(n)}
    return len(roots), max(dsu.size[r] for r in roots)

print(clubs(6, [(0, 1), (1, 2), (4, 5)]))  # (3, 3)

Problem 2. Finding useless cables

You get an ordered log of cables plugged between n office computers. A cable between two computers that are already connected only forms a loop and is useless. Print the indices of all useless cables in log order.

Approach: this is cycle detection in an undirected graph. When union returns False, both endpoints were already in the same set, so that edge closes a cycle.

def useless_cables(n, cables):
    dsu = DSU(n)
    return [i for i, (a, b) in enumerate(cables) if not dsu.union(a, b)]

print(useless_cables(4, [(0, 1), (1, 2), (2, 0), (2, 3), (3, 1)]))  # [2, 4]

Problem 3. Merging accounts

Each account has a name and a list of email addresses. Accounts that share at least one email belong to the same person, and this relation is transitive. Merge each person's accounts and produce a sorted email list per person.

Approach: use account indices as elements. Remember the first account that used each email; when the email shows up again, merge the two accounts. Finally collect emails by root.

from collections import defaultdict

def merge_accounts(accounts):
    dsu = DSU(len(accounts))
    owner = {}
    for i, (_, emails) in enumerate(accounts):
        for e in emails:
            if e in owner:
                dsu.union(i, owner[e])
            else:
                owner[e] = i
    groups = defaultdict(set)
    for e, i in owner.items():
        groups[dsu.find(i)].add(e)
    return [(accounts[r][0], sorted(es)) for r, es in groups.items()]

accounts = [
    ("kim", ["a@x.io", "b@x.io"]),
    ("kim", ["c@x.io"]),
    ("kim", ["b@x.io", "d@x.io"]),
]
print(merge_accounts(accounts))
# [('kim', ['a@x.io', 'b@x.io', 'd@x.io']), ('kim', ['c@x.io'])]