201. Implement Disjoint-Set Union

Why this mechanism works

Represent each component as a parent-pointer tree. A root points to itself. find shortens future routes by redirecting each visited node toward its grandparent. union first obtains roots, then attaches the smaller tree beneath the larger so depth grows slowly even before compression.

State and boundaries

parent links remain within a component and lead to its root. size is meaningful at roots; count equals the number of roots.

Protocol

1. Initialize each element as a singleton.
2. During find repeatedly halve the parent path.
3. During union find both roots.
4. Return false if equal.
5. Attach the smaller root, add its size, decrement count and return true.

Visualization

0<-1   2<-3
union(1,3): roots 0 and 2 merge
0<-2<-3; a later find(3) shortens its route to 0

Python implementation

from structures import run_operations

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

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

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a == b:
            return False
        if self.size[a] < self.size[b]:
            a, b = b, a
        self.parent[b] = a
        self.size[a] += self.size[b]
        self.count -= 1
        return True

Worked execution

Recorded state changes after statements; L numbers refer to the solution file. Long traces are explicitly condensed.

L5 __init__: self.parent = list(range(n))
self=DSU{'parent':[..., ..., ..., ...]}; n=4

L6 __init__: self.size = [1] * n
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...]}; n=4

L7 __init__: self.count = n
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':4}; n=4

L10 find: while x != self.parent[x]:
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':4}; x=0

L10 find: while x != self.parent[x]:
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':4}; x=1

L16 union: a, b = self.find(a), self.find(b)
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':4}; a=0; b=1

L23 union: self.count -= 1
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':3}; a=0; b=1

L10 find: while x != self.parent[x]:
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':3}; x=2

... 11 intermediate state changes omitted ...

L10 find: while x != self.parent[x]:
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':1}; x=0

L10 find: while x != self.parent[x]:
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':1}; x=3

L12 find: x = self.parent[x]
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':1}; x=0

L16 union: a, b = self.find(a), self.find(b)
self=DSU{'parent':[..., ..., ..., ...], 'size':[..., ..., ..., ...], 'count':1}; a=0; b=0

L5 __init__: self.parent = list(range(n))
self=DSU{'parent':[...]}; n=1

L6 __init__: self.size = [1] * n
self=DSU{'parent':[...], 'size':[...]}; n=1

L7 __init__: self.count = n
self=DSU{'parent':[...], 'size':[...], 'count':1}; n=1

L10 find: while x != self.parent[x]:
self=DSU{'parent':[...], 'size':[...], 'count':1}; x=0

L10 find: while x != self.parent[x]:
self=DSU{'parent':[...], 'size':[...], 'count':1}; x=0

L10 find: while x != self.parent[x]:
self=DSU{'parent':[...], 'size':[...], 'count':1}; x=0

L16 union: a, b = self.find(a), self.find(b)
self=DSU{'parent':[...], 'size':[...], 'count':1}; a=0; b=0

Why the solution is correct

Path shortening preserves the same reachable root and therefore component identity. Linking two distinct roots merges exactly their two sets without affecting others. Root sizes add because the sets were disjoint, and one fewer root justifies decrementing count. Size-based attachment and compression preserve a forest.

Time and space

O(n) initialization and O(alpha(n)) amortized time per operation over a sequence, with O(n) storage. alpha is the inverse Ackermann function; an individual operation is not guaranteed strict O(1).

Common mistakes

Attaching arbitrary nonroot nodes can disconnect represented relationships. Decrementing count after an already-connected union undercounts components.

Reasoning card

Transfer exercise

Support rollback by recording parent/size changes. Explain why ordinary path compression complicates undo.