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)        # n nodes, so n separate sets to start

    # Invariant: dsu.count = components among the edges seen so far.
    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}.

Reading the figure. Arrows point from a node to its parent. Green nodes are roots, and each root names one set. Blue nodes point to a parent. The thick green arrow is the link made in that step. Count the green nodes in any frame and you get count.

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.
    """
    # n - 1: a tree on n nodes has exactly one fewer edge than nodes.
    if len(edges) != n - 1:
        return False

    dsu = DisjointSet(n)        # n nodes, each its own set
    # union is False on a cycle edge. With n - 1 edges and no cycle,
    # the graph must be connected.
    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.
    # n nodes and n edges here, so len(edges) + 1 gives slots 0..n.
    dsu = DisjointSet(len(edges) + 1)

    # Invariant: the edges before this one form a forest with no cycle.
    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]]:

Reading the figure. Each box adds one input edge. A green edge joined two different sets. Gray edges were added in earlier steps. The dashed red edge connects 2 and 3, which already share root 1. union returns False there, and that edge is the answer.

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))  # one set per account to start
    owner_of: dict[str, int] = {}     # email -> index of the first account seen with it

    for index, account in enumerate(accounts):
        # account[1:]: skip index 0, which is the name. The rest are emails.
        for email in account[1:]:
            if email in owner_of:
                # A shared email means the same person: join the two accounts.
                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.
    # [0]: the name is the first entry of an account.
    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.

Reading the figure. Boxes on the left are accounts. Boxes in the middle are emails. Gray lines are emails seen for the first time, so they just record an owner. The amber line is x@z showing up a second time, which triggers union(2, 0). Same color on the left means same root, and the right side shows the merged output.

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

# The four neighbours of a cell: down, up, right, left.
# Each pair is (row change, column change). 1 moves forward, -1 back, 0 stays.
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)          # m * n: one slot per grid cell
    dsu.count = 0                     # the grid starts as all water

    is_land = [[False] * n for _ in range(m)]   # m rows of n cells, all water
    answer: list[int] = []

    # Invariant: dsu.count = number of islands among the land cells so far.
    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   # the neighbour cell
            # 0 <= ... < m and 0 <= ... < n keep the neighbour inside the grid.
            if 0 <= nr < m and 0 <= nc < n and is_land[nr][nc]:
                # Each successful merge takes the count back down by one.
                # r * n + c flattens (row, col) to one id: n cells per row.
                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.

Reading the figure. Each grid is the state after one addition. The cell with the amber border is the one just added. Cells of the same color belong to the same island. In frame 2 the new cell touches (0,0). One union succeeds, so the count goes up 1 and back down 1. In frame 3 the new cell only touches (0,1) at a corner, so it starts a new island.

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.