{"id":"dsa-graphs-union-find","slug":"union-find","title":"Challenge: union-find","text":"Think of a jigsaw puzzle on a table. Pieces snap together one pair at a time, and the loose pieces and finished chunks on the table are the connected components of a graph that only ever grows. In between snaps you may want to know whether two pieces already belong to the same chunk, or how many chunks are left. The DFS of step 3 can tell you, but it starts from nothing every time: O(V + E) for each question, again and again as the puzzle grows. **Union-find**, also called a **disjoint set union** (DSU), is built for groups that only merge. It answers \"same group?\" and joins two groups, each in almost O(1).\n\n### A parent list\n\nEach group is a small tree, kept in a list `parent`: `parent[x]` is the vertex above x. The vertex at the top of a group, its **root**, is its own parent, and it names the group. At the start, every vertex is alone, so `parent = list(range(n))`. Two operations do the work:\n\n- **find(x)** follows the parents up from x to the root. Two vertices are in the same group exactly when they have the same root.\n- **union(a, b)** finds both roots. If they differ, it hangs one root under the other, and the two groups are one.\n\n```python\ndef find(parent, x):\n    while parent[x] != x:\n        x = parent[x]\n    return x\n\n\ndef union(parent, a, b):\n    root_a, root_b = find(parent, a), find(parent, b)\n    if root_a != root_b:\n        parent[root_a] = root_b\n\n\nparent = list(range(7))\nunion(parent, 5, 2)\nunion(parent, 0, 6)\nunion(parent, 5, 0)\nprint(parent)\nprint(find(parent, 5) == find(parent, 6), find(parent, 1) == find(parent, 3))\n```\n\n```output\n[6, 1, 6, 3, 4, 2, 6]\nTrue False\n```\n\nAfter the three unions, 0, 2, 5 and 6 share the root 6, while 1, 3 and 4 are still alone. The tree needn't look like the edges you added: the last union named 5 and 0, but it hung 5's root, 2, under 0's root, 6, so `parent[5]` is still 2 and 5 is two steps below the root.\n\n### The trap: tall trees\n\nNothing here keeps the trees short. When person after person joins the group of person 0 with `union(0, i)`, each union hangs the old root under i, and 0 sinks to the bottom of a chain:\n\n```python\ndef find(parent, x):\n    while parent[x] != x:\n        x = parent[x]\n    return x\n\n\ndef depth(parent, x):\n    steps = 0\n    while parent[x] != x:\n        x = parent[x]\n        steps += 1\n    return steps\n\n\nfor n in [10, 100, 1000]:\n    parent = list(range(n))\n    for i in range(1, n):\n        root_a, root_b = find(parent, 0), find(parent, i)\n        if root_a != root_b:\n            parent[root_a] = root_b\n    print(n, \"people: 0 is\", depth(parent, 0), \"steps below its root\")\n```\n\n```output\n10 people: 0 is 9 steps below its root\n100 people: 0 is 99 steps below its root\n1000 people: 0 is 999 steps below its root\n```\n\nEach `find(0)` walks the whole chain so far, so the n unions take 1 + 2 + ... + (n - 1) steps: O(n²), 5 billion steps for 100,000 people.\n\n### Two fixes\n\n- **Union by size:** keep the size of each group at its root, and hang the root of the smaller group under the root of the bigger one. A vertex gets one step further from its root only when its group at least doubles in size, which can happen at most log₂ n times. So no tree is ever taller than log₂ n.\n- **Path compression:** once `find` has found the root, it walks the path again and points every vertex on it straight at the root. The next `find` from any of them takes one step.\n\n```python\ndef find(parent, x):\n    root = x\n    while parent[root] != root:\n        root = parent[root]\n    while parent[x] != root:\n        next_x = parent[x]\n        parent[x] = root\n        x = next_x\n    return root\n\n\ndef union(parent, size, a, b):\n    root_a, root_b = find(parent, a), find(parent, b)\n    if root_a == root_b:\n        return False\n    if size[root_a] < size[root_b]:\n        root_a, root_b = root_b, root_a\n    parent[root_b] = root_a\n    size[root_a] += size[root_b]\n    return True\n\n\nparent = list(range(6))\nsize = [1] * 6\ngroups = 6\nfor a, b in [(0, 1), (3, 4), (1, 2), (2, 0), (1, 4)]:\n    if union(parent, size, a, b):\n        groups -= 1\n    else:\n        print(\"redundant:\", a, b)\nprint(groups, \"groups\")\nprint(size[find(parent, 3)], find(parent, 0) == find(parent, 4))\n\nparent = list(range(1000))\nsize = [1] * 1000\nfor i in range(1, 1000):\n    union(parent, size, 0, i)\nprint(max(parent))\n```\n\n```output\nredundant: 2 0\n2 groups\n5 True\n0\n```\n\n`size[x]` only counts while x is a root, so the size of a group is `size[find(parent, x)]`. `union` returns `True` only when it joins two groups, so the number of groups starts at n and goes down by one with every `True`. In the second part, where 999 people join person 0's group one at a time, the big group's root 0 stays on top, and every other person hangs right under it: the largest parent is 0.\n\nWith both fixes, m operations on n vertices take O(m·α(n)) time, where α is a function that grows so slowly that it's at most 4 for any n that fits in a computer. In practice that's O(1) for each operation, with O(n) memory. Write `find` with loops, as here: a recursive `find` is shorter, but on a tall tree it hits Python's recursion limit before it has compressed anything.\n\n### Redundant edges\n\n`union` returns `False` when a and b are already in the same group. Then the edge between them connects nothing new: it is **redundant**, and in graph words it closes a cycle. That's how union-find finds the extra cable in a network, and how Kruskal's algorithm builds the cheapest network that connects every vertex: it goes through the edges from the cheapest up and skips each one that `union` says is redundant.\n","task":"Fill in the blanks: path compression's pointer, the comparison that picks the bigger group, and the root that the smaller root hangs under. The program should print:\n\n```\n0 [0, 0, 0, 0, 0, 5]\nTrue False\n[0, 0, 0, 0, 0, 0] 6\n```\n","kind":"fill","seed":"def find(parent, x):\n    root = x\n    while parent[root] != root:\n        root = parent[root]\n    while x != root:\n        next_x = parent[x]\n        parent[x] = [[root]]\n        x = next_x\n    return root\n\n\ndef union(parent, size, a, b):\n    root_a, root_b = find(parent, a), find(parent, b)\n    if root_a == root_b:\n        return False\n    if size[root_a] [[<]] size[root_b]:\n        root_a, root_b = root_b, root_a\n    parent[root_b] = [[root_a]]\n    size[root_a] += size[root_b]\n    return True\n\n\nparent = [0, 0, 1, 2, 3, 5]\nsize = [5, 1, 1, 1, 1, 1]\nprint(find(parent, 4), parent)\nprint(union(parent, size, 5, 2), union(parent, size, 3, 1))\nprint(parent, size[0])\n","solution":"def find(parent, x):\n    root = x\n    while parent[root] != root:\n        root = parent[root]\n    while x != root:\n        next_x = parent[x]\n        parent[x] = root\n        x = next_x\n    return root\n\n\ndef union(parent, size, a, b):\n    root_a, root_b = find(parent, a), find(parent, b)\n    if root_a == root_b:\n        return False\n    if size[root_a] < size[root_b]:\n        root_a, root_b = root_b, root_a\n    parent[root_b] = root_a\n    size[root_a] += size[root_b]\n    return True\n\n\nparent = [0, 0, 1, 2, 3, 5]\nsize = [5, 1, 1, 1, 1, 1]\nprint(find(parent, 4), parent)\nprint(union(parent, size, 5, 2), union(parent, size, 3, 1))\nprint(parent, size[0])\n","hints":["Path compression points every vertex on the path straight at the root that the first loop found.","After the `if`, `root_a` must be the root of the bigger group, so swap when a's group is smaller: `<`. Then the smaller root, `root_b`, hangs under `root_a`."],"tests":[{"name":"Output","stdin":"","expect":"0 [0, 0, 0, 0, 0, 5]\nTrue False\n[0, 0, 0, 0, 0, 0] 6","hidden":false}],"require":[],"forbid":[],"more":[{"task":"Write a class `UnionFind` for the vertices 0 to n - 1, with union by size and path compression:\n\n- `UnionFind(n)` starts with every vertex in a group of its own;\n- `find(x)` returns the root of x's group;\n- `union(a, b)` joins the groups of a and b and returns `True`, or returns `False` when they already were one group;\n- `size_of(x)` returns the number of vertices in x's group;\n- the attribute `groups` is the number of groups.\n\nThe hidden tests make 100,000 unions in which everyone joins one group: without the two fixes, a chain forms, and each `find` walks all of it.\n\nFor example, this program:\n\n```python\nsets = UnionFind(8)\nprint(sets.union(1, 5), sets.union(7, 5), sets.union(5, 1), sets.union(3, 0))\nprint(sets.groups, sets.size_of(7), sets.find(1) == sets.find(7), sets.find(0) == sets.find(5))\n```\n\nprints:\n\n```\nTrue True False True\n5 3 True False\n```\n","kind":"code","seed":"class UnionFind:\n    def __init__(self, n):\n        pass\n\n\nif __name__ == \"__main__\":\n    sets = UnionFind(8)\n    print(sets.union(1, 5), sets.union(7, 5), sets.union(5, 1), sets.union(3, 0))\n    print(sets.groups, sets.size_of(7), sets.find(1) == sets.find(7), sets.find(0) == sets.find(5))\n","solution":"class UnionFind:\n    def __init__(self, n):\n        self.parent = list(range(n))\n        self.size = [1] * n\n        self.groups = n\n\n    def find(self, x):\n        root = x\n        while self.parent[root] != root:\n            root = self.parent[root]\n        while self.parent[x] != root:\n            next_x = self.parent[x]\n            self.parent[x] = root\n            x = next_x\n        return root\n\n    def union(self, a, b):\n        root_a, root_b = self.find(a), self.find(b)\n        if root_a == root_b:\n            return False\n        if self.size[root_a] < self.size[root_b]:\n            root_a, root_b = root_b, root_a\n        self.parent[root_b] = root_a\n        self.size[root_a] += self.size[root_b]\n        self.groups -= 1\n        return True\n\n    def size_of(self, x):\n        return self.size[self.find(x)]\n\n\nif __name__ == \"__main__\":\n    sets = UnionFind(8)\n    print(sets.union(1, 5), sets.union(7, 5), sets.union(5, 1), sets.union(3, 0))\n    print(sets.groups, sets.size_of(7), sets.find(1) == sets.find(7), sets.find(0) == sets.find(5))\n","hints":["`__init__` makes three attributes: `self.parent = list(range(n))`, `self.size = [1] * n` and `self.groups = n`. `find` and `union` are the lesson's functions turned into methods, which use `self.parent` and `self.size` in place of the parameters.","`union` swaps the two roots when a's group is the smaller one, hangs `root_b` under `root_a`, adds the sizes at the new root and takes 1 off `groups`. It returns `False` early when the roots are the same.","The size of a group is kept at its root, so `size_of(x)` is `self.size[self.find(x)]`. Write `find` with two loops: one up to the root, one that points the path at it."],"tests":[{"name":"Test 1","stdin":"","call":"sets = UnionFind(8)\nprint(sets.union(1, 5), sets.union(7, 5), sets.union(5, 1), sets.union(3, 0))\nprint(sets.groups, sets.size_of(7), sets.find(1) == sets.find(7), sets.find(0) == sets.find(5))\n","expect":"True True False True\n5 3 True False","hidden":false},{"name":"Test 2","stdin":"","call":"sets = UnionFind(1)\nprint(sets.groups, sets.find(0), sets.size_of(0), sets.union(0, 0))\n","expect":"1 0 1 False","hidden":true},{"name":"Test 3","stdin":"","call":"sets = UnionFind(5)\nprint(sets.union(3, 4), sets.union(4, 3), sets.groups)\nprint(sets.union(0, 1), sets.union(1, 2), sets.size_of(0), sets.size_of(4))\nprint(sets.union(2, 4), sets.groups, sets.size_of(3), sets.find(0) == sets.find(4))\n","expect":"True False 4\nTrue True 3 2\nTrue 1 5 True","hidden":true},{"name":"Test 4","stdin":"","call":"sets = UnionFind(100_000)\nfor i in range(1, 100_000):\n    sets.union(0, i)\nprint(sets.groups, sets.size_of(0), sets.find(99_999) == sets.find(0))\n","expect":"1 100000 True","hidden":true,"big":true},{"name":"Test 5","stdin":"","call":"sets = UnionFind(100_000)\nfor i in range(1, 100_000):\n    sets.union(i, 0)\nprint(sets.groups, sets.size_of(99_999), sets.find(5) == sets.find(0))\n","expect":"1 100000 True","hidden":true,"big":true}],"require":[],"forbid":[{"pattern":"setrecursionlimit","message":"Don't raise the recursion limit: the browser's Python can crash. Write `find` with loops instead."}]},{"task":"A network of n computers, numbered 0 to n - 1, is built one cable at a time, in the order of the list `cables`. Each cable `(a, b)` connects two computers. Write `first_redundant(n, cables)` that returns the first cable whose two computers were already connected, directly or through other computers, by the cables before it, or `None` when every cable connects something new.\n\nA search from a to b before each cable would answer, but costs O(n) for each cable, and the hidden tests have 100,000 computers. Use union-find.\n\nFor example, in this program 2 and 0 are already connected through 1 when the cable `(2, 0)` comes:\n\n```python\nprint(first_redundant(5, [(0, 1), (1, 2), (3, 4), (2, 0), (1, 3)]))\nprint(first_redundant(3, [(0, 1), (1, 2)]))\n```\n\nprints:\n\n```\n(2, 0)\nNone\n```\n","kind":"code","seed":"def first_redundant(n, cables):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(first_redundant(5, [(0, 1), (1, 2), (3, 4), (2, 0), (1, 3)]))\n    print(first_redundant(3, [(0, 1), (1, 2)]))\n","solution":"def find(parent, x):\n    root = x\n    while parent[root] != root:\n        root = parent[root]\n    while parent[x] != root:\n        next_x = parent[x]\n        parent[x] = root\n        x = next_x\n    return root\n\n\ndef first_redundant(n, cables):\n    parent = list(range(n))\n    size = [1] * n\n    for a, b in cables:\n        root_a, root_b = find(parent, a), find(parent, b)\n        if root_a == root_b:\n            return (a, b)\n        if size[root_a] < size[root_b]:\n            root_a, root_b = root_b, root_a\n        parent[root_b] = root_a\n        size[root_a] += size[root_b]\n    return None\n\n\nif __name__ == \"__main__\":\n    print(first_redundant(5, [(0, 1), (1, 2), (3, 4), (2, 0), (1, 3)]))\n    print(first_redundant(3, [(0, 1), (1, 2)]))\n","hints":["Start with `parent = list(range(n))` and `size = [1] * n`, and go through the cables in order.","For each cable, find the roots of a and b. The same root means a and b are already connected: return the cable as a tuple.","Otherwise join the two groups by size, the smaller root under the bigger one. If no cable was redundant, return `None`."],"tests":[{"name":"Test 1","stdin":"","call":"print(first_redundant(5, [(0, 1), (1, 2), (3, 4), (2, 0), (1, 3)]))\nprint(first_redundant(3, [(0, 1), (1, 2)]))\n","expect":"(2, 0)\nNone","hidden":false},{"name":"first_redundant(2, [(0, 1), (1, 0)])","stdin":"","call":"first_redundant(2, [(0, 1), (1, 0)])\n","expect":"(1, 0)","hidden":true,"returns":true},{"name":"first_redundant(5, [(3, 4), (0, 2), (4, 0), (2, 3), (1, 2)])","stdin":"","call":"first_redundant(5, [(3, 4), (0, 2), (4, 0), (2, 3), (1, 2)])\n","expect":"(2, 3)","hidden":true,"returns":true},{"name":"print(first_redundant(1, []))","stdin":"","call":"print(first_redundant(1, []))\n","expect":"None","hidden":true},{"name":"first_redundant(6, [(4, 5), (3, 4), (0, 1), (1, 2), (2, 3), (5, 1), (0, 5)])","stdin":"","call":"first_redundant(6, [(4, 5), (3, 4), (0, 1), (1, 2), (2, 3), (5, 1), (0, 5)])\n","expect":"(5, 1)","hidden":true,"returns":true},{"name":"first_redundant(100_000, [(i, i + 1) for i in range(99_999)] + [(0, 99_999)])","stdin":"","call":"first_redundant(100_000, [(i, i + 1) for i in range(99_999)] + [(0, 99_999)])\n","expect":"(0, 99999)","hidden":true,"big":true,"returns":true},{"name":"first_redundant(100_000, [(0, i) for i in range(1, 100_000)] + [(500, 99_000)])","stdin":"","call":"first_redundant(100_000, [(0, i) for i in range(1, 100_000)] + [(500, 99_000)])\n","expect":"(500, 99000)","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[{"pattern":"setrecursionlimit","message":"Don't raise the recursion limit: the browser's Python can crash. Write `find` with loops instead."}]}]}