Union-find

Code in:JavaC#PythonJavaScript

In short: Union-find (disjoint set union) keeps groups of elements and tells whether two elements are in the same group; with path compression and union by size every operation is almost O(1).

Problem

One customer, three profiles

The shop has 200,000 customer profiles. The same customer signed up three times: once with a phone, once with her work email, once through the app. She gets three newsletters, her loyalty points are split three ways, and the report says you have more customers than you do.

The rule is simple: two profiles with a shared email or a shared phone are one customer. But it is contagious. If profile a shares a phone with b, and b shares an email with c, then a, b and c are one customer, even though a and c have nothing in common.

How do you find all the groups without comparing every profile with every other?

The idea

The first idea: a group label

Every profile carries a label: the number of its group. At first every profile is its own group. When two profiles turn out to be linked, the whole group of the second gets the label of the first.

class Relabel:
    def __init__(self, n):
        self._label = list(range(n))
        self.count = n

    def find(self, x):
        return self._label[x]

    def union(self, a, b):
        old, new = self._label[b], self._label[a]
        if old == new:
            return False
        for i, group in enumerate(self._label):
            if group == old:
                self._label[i] = new
        self.count -= 1
        return True

The program walks through the profiles and remembers, for every contact, which profile had it first. When it sees the contact again, it merges the two profiles:

profiles = [
    ["anna@mail.com", "555-0100"],
    ["555-0100", "a.smith@work.com"],
    ["peter@mail.com"],
    ["a.smith@work.com"],
    ["555-0300", "peter@mail.com"],
    ["nick@mail.com"],
    ["555-0300"],
    ["nick@mail.com", "555-0500"],
    ["anna@mail.com", "555-0300"],
]
accounts = Relabel(len(profiles))
owner = {}
for p, contacts in enumerate(profiles):
    for contact in contacts:
        if contact in owner:
            accounts.union(owner[contact], p)
        else:
            owner[contact] = p
print("Groups:", accounts.count)
print("Main:", *(accounts.find(p) for p in range(len(profiles))))

Output

Groups: 2
Main: 2 2 2 2 2 5 2 5 2

It works, but every merge walks over all profiles to find which ones to relabel. With n profiles and about n merges that is about n² looks: tens of billions for 200,000 profiles.

The idea

A forest of parents

Instead of every profile carrying the group label, every profile points to a parent: another profile in the group. The group is a tree, and the profile that points to itself is the root: the main profile of the group.

• find(x) climbs the parents until it reaches a root. Two profiles are in one group if find gives the same root for both.
• union(a, b) finds the roots of both and links one under the other. It touches one cell of the array, however big the group is.

On its own the forest can turn into a long chain, and then find is slow. So there are two improvements, and together they are union-find:

1. Union by size. Every root knows how many profiles its tree has. The smaller tree goes under the larger one, so depth grows slowly.
2. Path compression. Once find reaches the root, it walks the path again and links every profile on it straight to the root. The next find for them is one step.

See it

A merge from the inside

class UnionFind:
    def __init__(self, n):
        self._parent = list(range(n))
        self._size = [1] * n
        self.count = n

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

    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]
        self.count -= 1
        return True

union returns False if the profiles are already in one group: then nothing changes and count stays.

The same program as above, with UnionFind instead of Relabel:

profiles = [
    ["anna@mail.com", "555-0100"],
    ["555-0100", "a.smith@work.com"],
    ["peter@mail.com"],
    ["a.smith@work.com"],
    ["555-0300", "peter@mail.com"],
    ["nick@mail.com"],
    ["555-0300"],
    ["nick@mail.com", "555-0500"],
    ["anna@mail.com", "555-0300"],
]
accounts = UnionFind(len(profiles))
owner = {}
for p, contacts in enumerate(profiles):
    for contact in contacts:
        if contact in owner:
            accounts.union(owner[contact], p)
        else:
            owner[contact] = p
print("Groups:", accounts.count)
print("Main:", *(accounts.find(p) for p in range(len(profiles))))

Output

Groups: 2
Main: 0 0 0 0 0 5 0 5 0

class UnionFind:
    def __init__(self, n):
        self._parent = list(range(n))
        self._size = [1] * n
        self.count = n

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

    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]
        self.count -= 1
        return True

Step by step, in words

1. Nine customer profiles, each its own group: _parent[i] = i, every profile is a root. Groups: 9.
2. Profile 1 has the phone of 0: union(0, 1). Different roots, so 1 is linked under 0. Then 3 has the email of 1: find(1) climbs to 0, and 3 goes under 0 too.
3. The same for the others: 4 under 2 (shared email), 6 under 2 (shared phone), 7 under 5. A union changes one cell of the array: the parent of a root.
4. Profile 8 has the email of 0: 8 goes under 0. The group of 0 now has four profiles.
5. 8 also has the phone of 4. find(4) reaches 2, find(8) reaches 0. The tree of 2 is smaller (3 against 4), so its root goes under 0: one cell changes, three profiles change group.
6. At the end the program calls find for every profile. find(4) climbs 4 → 2 → 0, then links 4 straight to 0: path compression. The same for 6.
7. Two groups, main profiles 0 and 5. The tree is flat: the next find(4) is one step.

The two groups are the same as with labels, only the numbers differ: here the main profile is the root, with labels it is whichever label survived.

Try it

Predict the output

Six profiles and six links. After every union the program records the group count:

uf = UnionFind(6)
links = [[0, 1], [1, 2], [0, 2], [3, 4], [4, 5], [2, 3]]
counts = []
for a, b in links:
    uf.union(a, b)
    counts.append(uf.count)
print(counts)

Answer

[5, 4, 4, 3, 2, 1]

[1, 1, 1, 1, 1, 1]

[5, 4, 3, 2, 1, 0]

Check

The idea

The recursive find

Compression is often written as recursion: find on the parent gives the root, and the profile is linked to it. Every profile on the path does the same on the way back.

def find(self, x):
    if self._parent[x] != x:
        self._parent[x] = self.find(self._parent[x])
    return self._parent[x]

Try it yourself:

def find(self, x):
    if self._parent[x] != x:
        self._parent[x] = 
    return self._parent[x]

Check

With union by size the trees are low: at most 17 levels for 200,000 profiles (log₂ 200,000 ≈ 17.6). So the recursion is safe.

More precisely, if you are curious

Without union by size a tree can become a chain of n profiles, and the recursive find makes n nested calls, stacked on top of each other, as in the lesson "Recursion". Python stops at 1000 nested calls by default and raises RecursionError. The loop in the lesson has no such problem for any shape of tree.

The idea

What it costs

Run the three ways. There are n profiles and n random links "these two share a contact". The count is:

• with labels: every look at a profile's label during relabelling;
• with the plain forest and with union-find: every step of find, including the look at the root itself.

• Relabelling the whole group O(n²)
• Parent forest, no improvements O(n²)
• Union-find: compression and size O(n)

Algorithm	n = 10	n = 1,000	n = 10,000	n = 100,000
Relabelling the whole group	90	843,000	≈ 84,300,000	≈ 8,430,000,000
Parent forest, no improvements	30	25,641	2,102,613	≈ 210,261,300
Union-find: compression and size	30	3,295	33,355	334,195

The code runs in your browser, in JavaScript. The number of operations is the same in every language, because the algorithm is the same. The time depends on the machine, so it is only for comparison. “≈” is an estimate from the curve: we do not run an n that big, so the page does not freeze.

class UnionFind:
    def __init__(self, n):
        self._parent = list(range(n))
        self._size = [1] * n
        self.count = n

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

    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]
        self.count -= 1
        return True

With 1000 profiles the labels take 843,000 looks, the plain forest 25,641 steps, and union-find 3,295. With 10,000 the plain forest already takes 2,102,613 steps: its trees become long chains. Union-find takes 33,355, about 3.3 steps per link, and with 100,000 profiles 334,195: the same per link.

This is an amortized cost (see amortized complexity): one find can be long, but after it the path is flat and the next ones are short. Overall the operations cost nearly O(1) each.

More precisely, if you are curious

More precisely: with union by size (or by rank, see below) and path compression, m operations on n elements cost O(m · α(n)). α is the inverse Ackermann function. It grows so slowly that for any n that fits in the universe it is at most 4. That is why in practice people say "nearly constant". The proof is Tarjan's (1975) and it is not short.

Instead of size you can keep a rank: an upper bound on the height of the tree. The root with the smaller rank goes under the other, and a rank grows only when the ranks are equal. The cost is the same. Size is handier because it also answers "how many profiles are in the group".

Try it

Find the bug

A colleague wrote the union. The group count comes out as 2, but there are three different main profiles:

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[b] = ra
    self._size[ra] += self._size[rb]
    self.count -= 1
    return True

profiles = [
    ["anna@mail.com", "555-0100"],
    ["555-0100", "a.smith@work.com"],
    ["peter@mail.com"],
    ["a.smith@work.com"],
    ["555-0300", "peter@mail.com"],
    ["nick@mail.com"],
    ["555-0300"],
    ["nick@mail.com", "555-0500"],
    ["anna@mail.com", "555-0300"],
]
accounts = UnionFind(len(profiles))
owner = {}
for p, contacts in enumerate(profiles):
    for contact in contacts:
        if contact in owner:
            accounts.union(owner[contact], p)
        else:
            owner[contact] = p
print("Groups:", accounts.count)
print("Main:", *(accounts.find(p) for p in range(len(profiles))))

Output

Groups: 2
Main: 0 0 2 0 2 5 2 5 0

Which line has the bug?

Pick the line with the bug.

1def union(self, a, b):2 ra, rb = self.find(a), self.find(b)3 if ra == rb:4 return False5 if self._size[ra] < self._size[rb]:6 ra, rb = rb, ra7 self._parent[b] = ra8 self._size[ra] += self._size[rb]9 self.count -= 110 return True

Check

The idea

Where else

Union-find answers one question: are these two in one group, when groups only grow. The same holds for:

• delivery areas: depots and roads between them. Two depots are in one area if there is a route, possibly through other depots;
• a computer network: which machines can reach each other after every new cable;
• Kruskal's algorithm for the cheapest network of roads, which you will see in the graphs module (algo-graphs). There union-find tells whether a new road closes a cycle.

One limit: union-find does not split. If a shared phone turns out to be an office phone and has merged two different customers, the union cannot be undone. So the links are kept separately, and the groups are rebuilt when a link is removed. That is also why a phone that appears in hundreds of profiles is usually not used for merging at all.

Exercise

Delivery areas

Solve it in your editor.

Areas. There are 10 depots (0 … 9) and roads between them: 0-1, 1-2, 3-4, 5-6, 6-7, 7-5, 8-9. Add two methods to the class: connected(a, b), which tells whether two depots are in one area, and size_of(x), which returns how many depots are in the area of x. Print the number of areas, whether 0 and 2 and whether 2 and 3 are in one area, and how many depots are in the area of 5. Then open a new road 2-3 and print the number of areas again and whether 0 and 4 are in one area:

Output

Areas: 4
0 and 2: yes
2 and 3: no
Area of 5: 3 depots
New road 2-3. Areas: 3
0 and 4: yes

More tasks

• Groups by name. Print every group of profiles as a list: the main profile and all profiles under it.
• The largest group. Keep the largest group size while merging, without walking all roots at the end.
• Union by rank. Replace size with rank and compare the number of find steps with those in Bench.

Test

Check yourself

1. Profiles a and b share an email, b and c share a phone, and d shares nothing with anyone. How many customers is that?

4: every profile is separate3: a and c share no contact2: a, b and c are one customer, d is another

2. What does union do when the two profiles are already in one group?

Nothing: find gives the same root and the group count staysIt lowers the group count by oneIt links the root under itself

3. Why does the root of the smaller tree go under the root of the larger one, and not the other way round?

It does not matter, it only looks nicerSo the oldest profile stays the main oneSo depth grows only when the sizes are equal and never exceeds log₂ n

4. Two profiles were merged by a shared phone, and it turns out to be an office phone. How do you split them?

With a reverse unionBuild the groups again from the links, without the wrong oneRemove the cell from the parent array

Check

Start the lesson over

Related lessons

• Builds on: Lambdas · Time complexity: counting steps · Big-O: how the work grows · Amortized complexity
• Continues in: Minimum spanning tree · Greedy algorithms and when they work
• Tasks on this topic: Data structures