Data Structures & Algorithms Cheatsheet

Union-Find

Use this Data Structures & Algorithms reference while you build software engineering projects, review code, or refresh the syntax you reach for most.

Complexity

Disjoint Set Union tracks which elements share a group under merges. With both optimizations every operation is effectively constant.

Variant	find	union	Note
naive parent chain	O(n)	O(n)	degenerates to a linked list
union by rank or size only	O(log n)	O(log n)	tree height stays logarithmic
path compression only	O(log n) amortized	O(log n)	flattens on the way up
both, the standard	O(α(n))	O(α(n))	α is the inverse Ackermann, ≤ 4 for any real n
build over m operations	O(m · α(n))	—	treat as O(m)
count components	O(1)	—	maintain a counter on union
space	O(n)	O(n)	two int arrays

α(n) is below 5 for n up to 2^65536, so in practice quote these as O(1) amortized and note the true bound is O(α(n)).

When To Reach For This

Connectivity queries mixed with edge additions. DSU handles "are these connected" under incremental merges. A BFS would have to re-run for each query.
Counting connected components as edges arrive.
Cycle detection in an undirected graph. An edge whose endpoints already share a root closes a cycle.
Kruskal's minimum spanning tree, where DSU is the safety check that keeps the forest acyclic.
Grouping by equivalence: accounts sharing an email, strings connected by allowed swaps, equations of the form a/b = k.
Offline problems processed in reverse, where deletions become unions when the timeline is flipped.

DSU cannot split a set back apart, and it cannot answer "is there a path avoiding vertex x". If the problem needs deletions, either process it offline in reverse or use a link-cut tree.

Implementation

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

    def find(self, x):
        """Iterative path compression, no recursion limit risk."""
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:            # second pass flattens
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a, b):
        """Returns False if already joined, which detects a cycle."""
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:        # union by size
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.components -= 1
        return True

    def connected(self, a, b):
        return self.find(a) == self.find(b)

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

A one-pass recursive find is shorter but risks a stack overflow on 10⁵ elements:

def find(self, x):
    if self.parent[x] != x:
        self.parent[x] = self.find(self.parent[x])   # compress on unwind
    return self.parent[x]

class DSU {
    private final int[] parent, size;
    int components;

    DSU(int n) {
        parent = new int[n];
        size = new int[n];
        components = n;
        for (int i = 0; i < n; i++) { parent[i] = i; size[i] = 1; }
    }

    int find(int x) {
        while (parent[x] != x) {
            parent[x] = parent[parent[x]];    // path halving, one pass
            x = parent[x];
        }
        return x;
    }

    boolean union(int a, int b) {
        int ra = find(a), rb = find(b);
        if (ra == rb) return false;
        if (size[ra] < size[rb]) { int t = ra; ra = rb; rb = t; }
        parent[rb] = ra;
        size[ra] += size[rb];
        components--;
        return true;
    }
}

Path halving (parent[x] = parent[parent[x]]) is a single-pass alternative to full compression with the same asymptotic bound and fewer memory writes.

Non-Integer Keys

Map arbitrary labels to indices, or use a dict-backed DSU.

class StringDSU:
    def __init__(self):
        self.parent = {}
        self.size = {}

    def add(self, x):
        if x not in self.parent:
            self.parent[x] = x
            self.size[x] = 1

    def find(self, x):
        self.add(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

For a grid, flatten the cell coordinate: id = r * cols + c. That keeps the fast array implementation.

Canonical Problems

def count_components(n, edges):
    dsu = DSU(n)
    for u, v in edges:
        dsu.union(u, v)
    return dsu.components

def valid_tree(n, edges):
    """A tree is connected with exactly n - 1 edges and no cycle."""
    if len(edges) != n - 1:
        return False
    dsu = DSU(n)
    for u, v in edges:
        if not dsu.union(u, v):
            return False               # a cycle
    return dsu.components == 1

def redundant_connection(edges):
    """The first edge that closes a cycle. n nodes, n edges, 1-indexed."""
    dsu = DSU(len(edges) + 1)
    for u, v in edges:
        if not dsu.union(u, v):
            return [u, v]
    return []

def num_islands_dsu(grid):
    """Grid connectivity via flattened ids."""
    if not grid:
        return 0
    rows, cols = len(grid), len(grid[0])
    dsu = DSU(rows * cols)
    land = 0
    for r in range(rows):
        for c in range(cols):
            if grid[r][c] != "1":
                continue
            land += 1
            for dr, dc in ((1, 0), (0, 1)):     # right and down suffice
                nr, nc = r + dr, c + dc
                if (nr < rows and nc < cols and grid[nr][nc] == "1"
                        and dsu.union(r * cols + c, nr * cols + nc)):
                    land -= 1                    # a merge removes one island
    return land

def accounts_merge(accounts):
    """Merge accounts sharing any email. Union emails to the first email."""
    dsu = StringDSU()
    owner = {}
    for acct in accounts:
        name, emails = acct[0], acct[1:]
        for e in emails:
            owner[e] = name
            dsu.union(emails[0], e)
    from collections import defaultdict
    groups = defaultdict(list)
    for e in owner:
        groups[dsu.find(e)].append(e)
    return [[owner[root]] + sorted(es) for root, es in groups.items()]

def longest_consecutive_dsu(nums):
    """Union x with x + 1 whenever both exist."""
    dsu = StringDSU()
    present = set(nums)
    for x in present:
        if x + 1 in present:
            dsu.union(x, x + 1)
    return max((dsu.size[dsu.find(x)] for x in present), default=0)

Kruskal's MST

Sort the edges by weight and add each one whose endpoints are not already connected. DSU is what makes the acyclicity check O(1).

def kruskal(n, edges):
    """edges: (weight, u, v). Returns (total weight, chosen edges) or
    (inf, []) if the graph is disconnected. O(E log E)."""
    edges = sorted(edges)              # copy, do not reorder the caller's list
    dsu = DSU(n)
    total, chosen = 0, []
    for w, u, v in edges:
        if dsu.union(u, v):
            total += w
            chosen.append((u, v, w))
            if len(chosen) == n - 1:
                break                  # a spanning tree has exactly n - 1
    if len(chosen) != n - 1:
        return float("inf"), []
    return total, chosen

def min_cost_connect_points(points):
    """Complete graph on points, Manhattan distance. O(n^2 log n)."""
    n = len(points)
    edges = []
    for i in range(n):
        for j in range(i + 1, n):
            d = (abs(points[i][0] - points[j][0])
                 + abs(points[i][1] - points[j][1]))
            edges.append((d, i, j))
    return kruskal(n, edges)[0]

Prim's algorithm is the alternative, O(E log V) with a heap and better on dense graphs. Kruskal is better on sparse graphs and when the edges arrive sorted.

Weighted And Rollback Variants

A weighted DSU stores a value relative to the root, which answers ratio and offset queries.

class WeightedDSU:
    """Tracks value[x] / value[root[x]] for equations like a / b = k."""
    def __init__(self):
        self.parent = {}
        self.ratio = {}         # ratio[x] = value[x] / value[parent[x]]

    def find(self, x):
        if x not in self.parent:
            self.parent[x] = x
            self.ratio[x] = 1.0
        if self.parent[x] != x:
            root = self.find(self.parent[x])
            self.ratio[x] *= self.ratio[self.parent[x]]
            self.parent[x] = root
        return self.parent[x]

    def union(self, a, b, value):
        """Record a / b = value."""
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            self.parent[ra] = rb
            self.ratio[ra] = value * self.ratio[b] / self.ratio[a]

    def query(self, a, b):
        if a not in self.parent or b not in self.parent:
            return -1.0
        if self.find(a) != self.find(b):
            return -1.0
        return self.ratio[a] / self.ratio[b]

Rollback DSU skips path compression (so the structure is undoable), keeps union by rank for the O(log n) bound, and pushes each merge onto a stack so it can be reversed. That is what offline dynamic-connectivity solutions use.

Gotchas

Union roots, not elements. parent[b] = a instead of parent[find(b)] = find(a) corrupts the structure and is the single most common DSU bug.
Both optimizations or neither claim. Path compression without union by size still gives O(log n) amortized, and the naive version is O(n). Only both together give O(α(n)).
1-indexed input. Size the arrays n + 1 when vertices are numbered from 1, or the last vertex is out of range.
find before comparing. dsu.parent[a] == dsu.parent[b] is not a connectivity test, only find(a) == find(b) is.
Recursive find overflows on 10⁵ elements in Python. Use the iterative form.
No deletions. DSU cannot un-merge. Process the timeline in reverse or use rollback DSU.
Component count must be maintained on union, not recomputed by scanning, or you lose the O(1) query.
Directed graphs. DSU models undirected connectivity only. Strong connectivity needs Tarjan or Kosaraju.