Union-Find: Disjoint Set Data Structure

Union-Find (also called Disjoint Set Union or DSU) is a data structure that tracks a collection of disjoint sets. It supports two operations: find (which set contains this element?) and union (merge two sets). With path compression and union by rank, both operations achieve nearly O(1) amortized time.

Union-Find powers Kruskal’s MST algorithm, detects cycles in graphs, finds connected components, and solves many partitioning problems.

Introduction

Union-Find (also called Disjoint Set Union or DSU) is a data structure that tracks a collection of disjoint sets, supporting two operations: find—which set contains this element?—and union—merge two sets together. What makes it work is that with two simple optimizations, both operations achieve amortized O(α(n)) time, where α(n) is the inverse Ackermann function that grows so slowly it is effectively constant for all practical inputs. This near-constant time makes Union-Find one of the most practical data structures for graph algorithms and connectivity problems.

The power of Union-Find comes from its simplicity and versatility. It forms the backbone of Kruskal’s minimum spanning tree algorithm, where cycle detection during edge addition requires efficient connectivity queries. It finds connected components in graphs, detects cycles, and solves partitioning problems in social networks and image processing. The data structure is also easy to implement correctly—many developers get it wrong without path compression or union by rank, leaving significant performance on the table.

This guide covers the core operations with path compression and union by rank, explains why these optimizations give near-constant time, and walks through practical applications like counting connected components, Kruskal’s MST, and the “most stones removed” problem. You’ll also see when to use union by size instead of union by rank, and how to implement find iteratively to avoid recursion depth issues.

Basic Implementation

class UnionFind:
    """
    Union-Find with path compression and union by rank.
    Nearly O(1) amortized for both operations.
    """

    def __init__(self, n):
        self.parent = list(range(n))  # Parent pointer
        self.rank = [0] * n  # Approximate tree depth

    def find(self, x):
        """
        Find root with path compression.
        Flattens tree along the way for future speedups.
        """
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # Recursive compression
        return self.parent[x]

    def union(self, x, y):
        """
        Union by rank: attach smaller tree under larger tree.
        Returns True if merged, False if already in same set.
        """
        root_x = self.find(x)
        root_y = self.find(y)

        if root_x == root_y:
            return False  # Already in same set

        # Union by rank: attach smaller depth tree under larger
        if self.rank[root_x] < self.rank[root_y]:
            root_x, root_y = root_y, root_x  # Ensure root_x has >= rank

        self.parent[root_y] = root_x

        if self.rank[root_x] == self.rank[root_y]:
            self.rank[root_x] += 1

        return True

    def connected(self, x, y):
        """Check if two elements are in the same set."""
        return self.find(x) == self.find(y)

    def count(self):
        """Count number of distinct sets."""
        return sum(1 for i in range(len(self.parent)) if self.parent[i] == i)

Optimized Operations

class OptimizedUnionFind:
    """
    Union-Find with:
    - Path compression (find)
    - Union by rank (union)
    - Size tracking
    """

    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.size = [1] * n  # Track component sizes

    def find(self, x):
        """Find with iterative path compression."""
        root = x
        while self.parent[root] != root:
            root = self.parent[root]

        # Compress path
        while self.parent[x] != root:
            next_x = self.parent[x]
            self.parent[x] = root
            x = next_x

        return root

    def union(self, x, y):
        """Union by size: attach smaller tree under larger."""
        root_x = self.find(x)
        root_y = self.find(y)

        if root_x == root_y:
            return False

        # Attach smaller size tree under larger
        if self.size[root_x] < self.size[root_y]:
            root_x, root_y = root_y, root_x

        self.parent[root_y] = root_x
        self.size[root_x] += self.size[root_y]

        return True

    def component_size(self, x):
        """Get size of component containing x."""
        return self.size[self.find(x)]

    def max_component_size(self):
        """Largest component size."""
        return max(self.size)

Common Applications

# Number of connected components in graph
def count_connected_components(n, edges):
    """Count components using union-find."""
    uf = UnionFind(n)
    for u, v in edges:
        uf.union(u, v)
    return uf.count()


# Detect cycle in undirected graph
def has_cycle(n, edges):
    """Detect if graph has a cycle."""
    uf = UnionFind(n)
    for u, v in edges:
        if not uf.union(u, v):
            return True  # Already connected = cycle
    return False


# Kruskal's MST using union-find
def kruskal_mst(n, edges):
    """
    MST using Kruskal's algorithm.
    edges: list of (weight, u, v)
    """
    # Sort edges by weight
    edges.sort(key=lambda x: x[0])

    uf = UnionFind(n)
    mst = []
    total_weight = 0

    for weight, u, v in edges:
        if uf.union(u, v):
            mst.append((u, v, weight))
            total_weight += weight

        if len(mst) == n - 1:  # MST has n-1 edges
            break

    return mst, total_weight


# Most stones removed (LeetCode 947)
def remove_stones(stones):
    """
    Maximum stones that can be removed = total stones - number of components.
    Use union-find on stone positions.
    """
    uf = UnionFind(20000)  # 0-9999 for row, 10000-19999 for col

    for x, y in stones:
        uf.union(x, 10000 + y)  # Connect row x to column y

    # Count distinct roots
    roots = set(uf.find(x) for x, _ in stones)
    return len(stones) - len(roots)


# Longest consecutive sequence (LeetCode 128)
def longest_consecutive(nums):
    """
    Find longest consecutive sequence length.
    Use union-find to group consecutive numbers.
    """
    if not nums:
        return 0

    uf = UnionFind(len(nums))
    num_to_idx = {num: i for i, num in enumerate(nums)}

    # Union each number with its neighbor if it exists
    for num in nums:
        if num - 1 in num_to_idx:
            uf.union(num_to_idx[num], num_to_idx[num - 1])
        if num + 1 in num_to_idx:
            uf.union(num_to_idx[num], num_to_idx[num + 1])

    # Find largest component size
    max_size = 0
    for i in range(len(nums)):
        if uf.find(i) == i:
            max_size = max(max_size, uf.size[i])

    return max_size

Time Complexity Analysis

Operation	Naive	With Path Compression	With Both Optimizations
find	O(n)	O(log n)	α(n) ≈ O(1)
union	O(n)	O(log n)	α(n) ≈ O(1)
connected	O(n)	O(log n)	α(n) ≈ O(1)

α(n) is the inverse Ackermann function, which grows so slowly that for all practical values of n, α(n) ≤ 4.

Why Path Compression and Union by Rank Work

# Without optimizations: tree can degenerate to linked list (O(n) operations)

# With path compression:
# find(1) → 1 → 3 → 5 (find root)
# After: parent[1]=5, parent[3]=5, parent[5]=5 (flattened)

# With union by rank:
# Always attach smaller tree under larger
# Tree height ≤ log(n) because each merge at least doubles size of one component

Union-Find With Path Compression and Rank

graph TD
    subgraph Before["Before Path Compression (find 1)"]
        direction TB
        N1["1"]
        N2["2"]
        N3["3"]
        N4["4"]
        N5["5"]

        N1 --> N2
        N2 --> N3
        N3 --> N4
        N4 --> N5

        N5Style["5 is root"]
    end

    subgraph After["After Path Compression"]
        direction TB
        A1["1"]
        A2["2"]
        A3["3"]
        A4["4"]
        A5["5"]

        A1 -.-> A5
        A2 -.-> A5
        A3 -.-> A5
        A4 -.-> A5
        A5["5 (root)"]
    end

    subgraph UnionByRank["Union by Rank Example"]
        direction LR
        R1["rank=2, size=4"]
        R1 --> R2["rank=1, size=2"]
        R1 --> R3["rank=1, size=1"]
        R1 --> R4["rank=0"]

        R5["Attach smaller under larger"]
    end

When to Use Union-Find

Problem	Union-Find Use
Kruskal’s MST	Detect cycles while adding edges
Connected components	Group related items
Cycle detection	Check if edge connects nodes in same component
Network connectivity	Is x connected to y?
Image processing	Connected components in pixels
Clustering	Group similar items

Production Failure Scenarios

1. Path compression skipped in recursive find: If you implement find recursively but forget the assignment self.parent[x] = self.find(self.parent[x]), path compression never happens. The tree degrades to O(n) height over time. Even iterative implementations can miss path compression if the compression loop doesn’t update all nodes along the chain. Every find must flatten—test this with a long chain before deployment.

2. Union by rank/size not applied consistently: If you sometimes union with arbitrary ordering instead of by rank or size, the tree height bound breaks. A sequence of unions in bad order produces a linear chain, turning near-O(1) operations into O(n). This is especially insidious because it works fine on small test cases but degrades on large inputs.

3. Rank overflow with large n: The rank field is typically a small integer. If you have millions of elements and keep merging trees of equal rank, rank can overflow an 8-bit or 16-bit integer. The tree still works but the rank comparison logic breaks, potentially causing incorrect unions.

4. Parent array not initialized for all elements: When using union-find in a dynamic context (adding elements after initialization), new elements need proper parent and rank initialization. If you skip this, find on uninitialized elements reads garbage from the arrays and corrupts the structure.

Observability Checklist

Track union-find operations to catch degradation and correctness issues.

Core Metrics

• find operation count and latency per call
• union operation count and merge rate
• Tree depth distribution across all roots
• Rank distribution across all elements
• Component count over time

Health Signals

• find latency increasing: tree height growing beyond log n
• Union merge rate dropping: all unions are between same-set elements (may indicate logic error)
• Rank values growing beyond log₂(n): union-by-rank broken
• Component count not decreasing as expected: union operations failing silently
• find returning unexpected values: parent array corrupted

Alerting Thresholds

• find latency > log₂(n) consistently: tree degradation
• Rank values exceeding 30 for n < 10⁶: rank tracking error
• union returning False rate > 50% without actual cycles: logic error
• Component count stable when it should be changing: union not executing

Distributed Tracing

• Trace each find with path length and nodes visited
• Include root node and compression ratio in span metadata
• Track union with merged component sizes
• Correlate slow finds with specific tree shapes (deep chains)

Security Notes

Union-find has specific security concerns when attacker controls operations.

Rank manipulation to cause worst-case performance: An attacker who controls union order can craft sequences that violate the union-by-rank guarantee, creating deep trees. While find is still O(log n), the constant factor matters under heavy load.

Fix: Use randomized union by shuffling operations if adversarial. Or use union by size which is harder to manipulate predictably.

Parent array corruption via overflow: If the parent array stores indices and an attacker can trigger index overflow (through large n values), the structure can be corrupted to point to arbitrary memory locations.

Fix: Validate n against maximum bounds before initialization. Use language-level bounds checking where available.

Cycle detection bypass via timing: In cycle detection contexts, if timing differences in find operations reveal tree structure, an attacker could use this to craft “safe” edges that actually create cycles but appear safe by timing.

Fix: Ensure find operations take constant time regardless of path length. Use branchless implementations to eliminate timing side channels.

Quick Recap Checklist

• Union-Find tracks disjoint sets efficiently
• find(x): returns root of set containing x
• union(x, y): merges sets containing x and y
• Path compression flattens tree on find
• Union by rank attaches smaller tree under larger
• Amortized O(α(n)) ≈ O(1) per operation
• Used in Kruskal’s MST, cycle detection, connected components

Trade-off Analysis

Approach	find Complexity	union Complexity	Space	Notes
Naive	O(n)	O(n)	O(n)	No optimizations
Path Compression Only	O(log n)	O(log n)	O(n)	Flattening on find
Union by Rank Only	O(log n)	O(log n)	O(n)	Height bound log n
Both Optimizations	α(n) ≈ O(1)	α(n) ≈ O(1)	O(n)	Inverse Ackermann
Union by Size	O(log n)	O(log n)	O(n)	Same as rank
Weighted Quick-Union	O(n) worst	O(log n)	O(n)	Without path compression

When to Prefer Each

• Path compression only: Good when unions are rare, finds are frequent
• Union by rank only: Good when finds and unions are balanced
• Both (standard): Best overall performance for most use cases
• Union by size: Better when you need component size tracking
• Weighted quick-union: Legacy approach, path compression superior

Interview Questions

Further Reading

• CLRS Chapter 21: Data Structures for Disjoint Sets — The canonical textbook treatment with formal proofs
• CP-Algorithms: Union-Find — Practical implementation variants and applications
• Visualgo: Union-Find — Interactive visualization of union-find operations
• LeetCode Union-Find Problems — Practice problems categorized by difficulty
• Kruskal’s Algorithm — Union-Find’s primary application in MST finding
• Ackermann Function — Understanding why α(n) is effectively constant

Conclusion

Union-Find gives you near-constant time set operations—find and union both run in O(α(n)), which is effectively O(1) for any practical input size. Path compression flattens the tree on each find, and union by rank keeps trees shallow. It’s the standard tool for cycle detection in Kruskal’s MST, grouping connected components, and any problem involving disjoint sets. For graph traversal basics, see Graphs: Adjacency List and Matrix, or Topological Sort for another graph ordering algorithm.