Union-Find (Disjoint Set Union)

1. Number of Connected Components Medium

Problem

Given n nodes labelled 0 to n-1 and a list of undirected edges, return the number of connected components.

Approach

• Start with n components, one per node. Every edge that joins two different components reduces the count by one. Edges inside a component change nothing.
• The count field does the whole job, so there is no final pass to count distinct roots.
• DFS also solves this in O(V + E) and is a fine answer here, because the edges are all known up front. Say that, then give the Union-Find version as the one that extends to streaming edges.

Solution

def count_components(n: int, edges: list[list[int]]) -> int:
    """Number of connected components in an undirected graph.

    Args:
        n: Node count. Nodes are labelled 0 .. n-1.
        edges: Undirected edges as [a, b] pairs.

    Returns:
        The number of connected components.

    Example:
        >>> count_components(5, [[0, 1], [1, 2], [3, 4]])
        2
    """
    dsu = DisjointSet(n)

    for a, b in edges:
        dsu.union(a, b)      # only a real merge decrements dsu.count

    return dsu.count

Walkthrough

n = 5, edges [[0,1], [1,2], [3,4]]:

Components are {0, 1, 2} and {3, 4}.

Effectively O(E). Quoting the α factor and then saying “which is at most 4, so effectively linear” is the complete answer.

The related one-liner

def valid_tree(n: int, edges: list[list[int]]) -> bool:
    """A graph is a tree when it is connected and has exactly n-1 edges.

    Equivalently: n-1 edges and no edge is ever redundant.
    """
    if len(edges) != n - 1:
        return False

    dsu = DisjointSet(n)
    return all(dsu.union(a, b) for a, b in edges)

If any union returns False, that edge closed a cycle, so it is not a tree. With exactly n - 1 edges and no cycle, connectivity follows automatically. That is a neat two-fact argument worth having ready.

Edge cases to raise

• No edges: the answer is n.
• All nodes in one chain: the answer is 1.
• Duplicate edges: the second one is a no-op, since the endpoints already share a root.
• Self-loop [2, 2]: also a no-op. Correct.
• n = 0: the answer is 0.

2. Redundant Connection Medium

Problem

A tree on n nodes had one extra edge added, creating exactly one cycle. Given the edge list, return the edge that can be removed to restore a tree. If several answers exist, return the one that appears last in the input.

Approach

• Process the edges in the order given. The first edge whose two endpoints are already connected is the one that closes the cycle.
• Because the graph has exactly one extra edge, that first failure is also the last such edge in the input. So the tie-break rule is satisfied by scanning forwards and returning immediately, with no extra logic.
• The bool from union is the entire solution. This problem exists to test whether you have that return value.
• Nodes are labelled 1..n, so size the structure at n + 1.

Solution

def find_redundant_connection(edges: list[list[int]]) -> list[int]:
    """The edge that closes the single cycle in a tree plus one extra edge.

    Args:
        edges: Undirected edges over nodes 1..n, in the order they were added.

    Returns:
        The edge to remove, the last one in the input that closes a cycle.

    Raises:
        ValueError: If no edge closes a cycle.

    Example:
        >>> find_redundant_connection([[1, 2], [1, 3], [2, 3]])
        [2, 3]
    """
    # Nodes are 1-indexed, so allocate one extra slot and ignore index 0.
    dsu = DisjointSet(len(edges) + 1)

    for a, b in edges:
        # union returns False when a and b already share a root, which
        # means this edge creates a cycle.
        if not dsu.union(a, b):
            return [a, b]

    raise ValueError("no redundant edge: the graph is already a tree")

Walkthrough

[[1,2], [1,3], [2,3]]:

Why the sizing is len(edges) + 1

A tree on n nodes has n - 1 edges, and this graph has one more, so len(edges) == n. Nodes run from 1 to n, so the arrays need n + 1 slots. Slot 0 is allocated and never touched. Getting this wrong gives an IndexError on the highest-numbered node, which is easy to miss on a small test.

Edge cases to raise

• The cycle is the very first pair of edges, such as [[1,2],[2,1]]: caught on the second edge.
• The extra edge is last: the common case, and it works because scanning forwards finds it last too.
• A self-loop: union(x, x) returns False immediately, so it is reported. Ask whether self-loops can appear.
• Several cycles: the problem promises exactly one. Say what you are relying on.

3. Accounts Merge Medium

Problem

Each account is a name followed by a list of emails. Two accounts belong to the same person if they share at least one email. Merge them, and return each merged account as the name followed by its emails in sorted order. The same name may belong to different people.

The modelling decision

Approach

• Keep a map from each email to the first account index that mentioned it.
• Scanning the emails of an account: if the email is new, record it as owned by this account. If it has been seen, union this account with the one that first claimed it.
• After the pass, walk the email map once and bucket every email under find(owner). Each bucket is one person.
• Sort each bucket and prepend the name, taken from any account in the group. The root index is convenient.

Solution

from collections import defaultdict


def accounts_merge(accounts: list[list[str]]) -> list[list[str]]:
    """Merge accounts that share any email address.

    Args:
        accounts: Each entry is [name, email1, email2, ...].

    Returns:
        One entry per person: [name, ...emails sorted ascending].

    Example:
        >>> accounts_merge([["A", "x@z"], ["B", "y@z"], ["A", "x@z", "w@z"]])
        [['A', 'w@z', 'x@z'], ['B', 'y@z']]
    """
    dsu = DisjointSet(len(accounts))
    owner_of: dict[str, int] = {}     # email -> index of the first account seen with it

    for index, account in enumerate(accounts):
        for email in account[1:]:
            if email in owner_of:
                dsu.union(index, owner_of[email])
            else:
                owner_of[email] = index

    # Bucket every email under the root of the account that owns it.
    groups: defaultdict[int, set[str]] = defaultdict(set)
    for email, index in owner_of.items():
        groups[dsu.find(index)].add(email)

    # Every account in a group carries the same name, so the root's will do.
    return [
        [accounts[root][0], *sorted(emails)]
        for root, emails in groups.items()
    ]

Walkthrough

Accounts [["A","x@z"], ["B","y@z"], ["A","x@z","w@z"]]:

Accounts 0 and 2 now share a root, so their emails bucket together. Account 1 stands alone even though nothing else was needed to separate it.

E is the total number of emails. The Union-Find work is effectively linear, so the sort dominates. Naming the sort as the bottleneck, rather than the merging, is the sharp answer here.

Why two accounts with the same name are not merged

Edge cases to raise

• An account with no emails: contributes nothing and disappears from the output. Ask whether it should appear as a bare name.
• A chain merge, where 1 shares with 2 and 2 shares with 3 but 1 and 3 share nothing: all three merge, via transitivity. This is the case that rules out simple pairwise comparison.
• Two people with the same name and no shared email: stay separate.
• The same email listed twice inside one account: harmless, since union(i, i) is a no-op.

4. Number of Islands II Hard

Problem

You start with an m × n grid of water. Land is added one cell at a time. After each addition, report the current number of islands. Return the list of counts.

Why this is the Union-Find problem

Approach

• Flatten the 2D grid to 1D: cell (r, c) is index r * n + c. That lets the standard integer class work unchanged, with no coordinate keys.
• Start count at zero, because there is no land yet. This is the one place the class needs adjusting.
• Adding land: increment count optimistically, treating the new cell as its own island. Then union it with each of the four neighbours that is already land. Every successful union decrements count, so the arithmetic settles itself.
• Guard against a repeated position. Adding land where land already exists must not change anything, and without the guard it would inflate the count.

Solution

DIRECTIONS: tuple[tuple[int, int], ...] = ((1, 0), (-1, 0), (0, 1), (0, -1))


def num_islands2(m: int, n: int, positions: list[list[int]]) -> list[int]:
    """Island count after each land addition to an initially empty grid.

    Args:
        m: Number of rows.
        n: Number of columns.
        positions: Cells [row, col] turned into land, in order.

    Returns:
        The island count after each addition.

    Example:
        >>> num_islands2(3, 3, [[0, 0], [0, 1], [1, 2], [2, 1]])
        [1, 1, 2, 3]
    """
    dsu = DisjointSet(m * n)
    dsu.count = 0                     # the grid starts as all water

    is_land = [[False] * n for _ in range(m)]
    answer: list[int] = []

    for r, c in positions:
        if is_land[r][c]:
            answer.append(dsu.count)  # repeat position: nothing changes
            continue

        is_land[r][c] = True
        dsu.count += 1                # a brand new island, for the moment

        for dr, dc in DIRECTIONS:
            nr, nc = r + dr, c + dc
            if 0 <= nr < m and 0 <= nc < n and is_land[nr][nc]:
                # Each successful merge takes the count back down by one.
                dsu.union(r * n + c, nr * n + nc)

        answer.append(dsu.count)

    return answer

Walkthrough

m = 3, n = 3, positions [[0,0], [0,1], [1,2], [2,1]]:

Notice [1,2] is diagonal to [0,1], and diagonals do not connect.

k is the number of additions. Each one does at most four unions, so the work per addition is constant.

The optimistic counting trick

Edge cases to raise

• A repeat position: guarded, and the previous count is repeated.
• A cell that bridges two or three existing islands: handled by the optimistic count.
• A 1×1 grid: one addition gives [1].
• Empty position list: returns [].
• Only the four cardinal directions connect. Confirm that diagonals do not.
• If m × n is huge but k is small, allocating the full grid may be wasteful. A dictionary keyed by coordinate avoids it, at some constant-factor cost. Worth a sentence.