Union-Find (DSU) & Trie

Two "if you know it, it's easy; if you don't, it's impossible" data structures. Each exists because a naive structure answers one query well and collapses on a different query the problem actually asks.

Why these two matter

A hashset tells you membership but cannot cheaply merge two evolving partitions. Union-Find answers "are these two things in the same group?" and "merge these two groups" — both in near-constant time. A hashset tells you "is this exact word present?" but knows nothing about prefixes; a Trie answers "which stored words share this prefix?" in time proportional to the prefix length, independent of how many words you stored.

Both structures replace "search everything every time" with "follow pointers a tiny number of times". Union-Find follows parent pointers up to a root; a Trie follows child pointers down a path. In both, the shape of the pointer graph is the optimization.

Part 1 — Union-Find (Disjoint Set Union, "DSU")

You have n elements, initially each in its own group. You receive a stream of operations: union(a, b) merges the groups containing a and b; find(a) returns a canonical id ("representative") for a's group; connected(a, b) asks whether they share a representative.

The naive idea — store a label per element, relabel one whole group on each union — costs O(n) per union, O(n²) overall. Too slow.

The forest idea

Represent each group as a tree. Every element points to a parent; the root (an element that points to itself) is the group's representative. To find the representative, walk parent pointers to the root. To union two groups, make one root point at the other. Both operations are O(tree height). The whole game is now: keep trees short. Two cheap tricks do it:

1. Union by rank/size — when merging, attach the smaller tree under the larger one. This keeps height logarithmic on its own.
2. Path compression — during find(), point every node we pass directly at the root. The next find() on any of them is O(1). This flattens trees over time.

class DSU:
    """
    Disjoint Set Union over integer ids 0..n-1, with union by size and path
    compression. `count` tracks the number of distinct groups.
    """

    def __init__(self, n: int) -> None:
        # parent[i] == i means i is a root. Initially every node is its own root.
        self.parent: List[int] = list(range(n))
        # size[i] is only meaningful when i is a root: how many nodes in its tree.
        self.size: List[int] = [1] * n
        # number of disjoint groups; starts at n, drops by 1 on each real union.
        self.count: int = n

    def find(self, x: int) -> int:
        """Return the root of x's tree, compressing the path on the way up."""
        root = x
        # Step 1: walk to the root.
        while self.parent[root] != root:
            root = self.parent[root]
        # Step 2: PATH COMPRESSION — relink every node on the path to root.
        while self.parent[x] != root:
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a: int, b: int) -> bool:
        """
        Merge the groups of a and b. Returns True if they were previously
        separate (a real merge happened), False if already in the same group.
        The False case is exactly a CYCLE signal for undirected graphs.
        """
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False  # already together → adding this edge forms a cycle
        # UNION BY SIZE: hang the smaller tree under the larger root.
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.count -= 1
        return True

    def connected(self, a: int, b: int) -> bool:
        """Same group? O(1) amortized."""
        return self.find(a) == self.find(b)

Problem 1 — number of connected components

Union every edge; the answer is simply DSU.count at the end. O(V + E·α).

def count_components(n: int, edges: List[List[int]]) -> int:
    dsu = DSU(n)
    for a, b in edges:
        dsu.union(a, b)
    return dsu.count

Problem 2 — redundant connection (LeetCode 684)

A tree on n nodes has exactly n-1 edges and no cycle. We're given n edges (1-indexed) forming exactly one extra edge that creates a cycle. Return the last edge that, when added, joins two already-connected nodes. The key insight: union() returning False is the cycle detector — for free.

def find_redundant_connection(edges: List[List[int]]) -> List[int]:
    dsu = DSU(len(edges) + 1)  # +1 because nodes are labeled 1..n
    for a, b in edges:
        if not dsu.union(a, b):  # already connected → this edge closes a cycle
            return [a, b]
    return []  # problem guarantees one exists; defensive default

Problem 3 — accounts merge (LeetCode 721)

Each account is [name, email1, email2, ...]. Two accounts belong to the same person iff they share any email. "Share any email" is a transitive equivalence that arrives incrementally as we scan accounts — classic union-find. Give each account an integer id (its index); map each email to the first account that owned it, and union when an email reappears.

def accounts_merge(accounts: List[List[str]]) -> List[List[str]]:
    dsu = DSU(len(accounts))
    email_owner: Dict[str, int] = {}   # email -> account id that first claimed it
    email_name: Dict[str, str] = {}    # email -> display name
    for i, acct in enumerate(accounts):
        name = acct[0]
        for email in acct[1:]:
            email_name[email] = name
            if email in email_owner:
                dsu.union(i, email_owner[email])  # same person across accounts
            else:
                email_owner[email] = i
    # Bucket every email under its group's representative account.
    groups: Dict[int, List[str]] = {}
    for email, owner in email_owner.items():
        root = dsu.find(owner)
        groups.setdefault(root, []).append(email)
    # Build the required [name, sorted emails...] rows.
    result: List[List[str]] = []
    for root, emails in groups.items():
        emails_sorted = sorted(set(emails))
        result.append([email_name[emails_sorted[0]]] + emails_sorted)
    return result

Problem 4 — number of islands via DSU (LeetCode 200)

Grid of '1' (land) and '0' (water); count maximal 4-connected land blobs. Flatten cell (r,c) to id r*cols + c. For each land cell, union it with its land neighbor to the left and above — those two directions suffice to connect the whole component as we scan left-to-right, top-down.

def num_islands(grid: List[List[str]]) -> int:
    if not grid or not grid[0]:
        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
            idx = r * cols + c
            # Union with left neighbor if it's land.
            if c > 0 and grid[r][c - 1] == "1":
                if dsu.union(idx, idx - 1):
                    land -= 1  # two land cells merged → one fewer island
            # Union with the neighbor above if it's land.
            if r > 0 and grid[r - 1][c] == "1":
                if dsu.union(idx, idx - cols):
                    land -= 1
    # `land` counted every land cell, then subtracted one for each successful
    # merge → exactly the number of connected land components.
    return land

DFS/BFS also works and is often shorter; DSU shines if islands keep merging as you add land dynamically — e.g. LeetCode 305, "Number of Islands II".

Part 2 — Trie (prefix tree)

Store a set of strings so that prefix questions are cheap: "is 'apple' stored?" (search), "is any word starting 'app'?" (startsWith), "list all words starting 'app'" (autocomplete). A hashset answers exact membership in O(L) but knows nothing about prefixes — to find words starting "app" it must scan every key.

The tree idea

Build a tree where each edge is one character and each root-to-node path spells a prefix. Words that share a prefix share that part of the path. A boolean flag on a node marks "a stored word ends exactly here".

    root
     └─a─p─p   (is_word=True at the second 'p' if "app" was inserted)
           └─l─e  (is_word=True → "apple")

class TrieNode:
    """
    One node = one position in some prefix. `children` maps a single character
    to the next node. `is_word` marks that a complete inserted word ends here.
    """

    __slots__ = ("children", "is_word")

    def __init__(self) -> None:
        self.children: Dict[str, "TrieNode"] = {}
        self.is_word: bool = False


class Trie:
    """Standard prefix tree: insert / search / startsWith."""

    def __init__(self) -> None:
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        """Walk/create the path char by char, then mark the end. O(L)."""
        node = self.root
        for ch in word:
            # setdefault creates the child only if this edge doesn't exist yet.
            node = node.children.setdefault(ch, TrieNode())
        node.is_word = True

    def _walk(self, prefix: str) -> Optional[TrieNode]:
        """Follow the path for `prefix`; return the node it ends at, or None."""
        node = self.root
        for ch in prefix:
            nxt = node.children.get(ch)
            if nxt is None:
                return None  # path breaks → prefix not present
            node = nxt
        return node

    def search(self, word: str) -> bool:
        """True iff the EXACT word was inserted (path exists AND ends a word)."""
        node = self._walk(word)
        return node is not None and node.is_word

    def starts_with(self, prefix: str) -> bool:
        """True iff some inserted word has this prefix (path exists)."""
        return self._walk(prefix) is not None

Wildcard search — WordDictionary (LeetCode 211)

addWord / search where the search pattern may contain '.' as a wildcard that matches any single character. The wildcard is why we need DFS: at a '.' we can't follow one edge, so we branch into all children and succeed if any branch matches the rest.

class WordDictionary:
    def __init__(self) -> None:
        self.root = TrieNode()

    def add_word(self, word: str) -> None:
        node = self.root
        for ch in word:
            node = node.children.setdefault(ch, TrieNode())
        node.is_word = True

    def search(self, word: str) -> bool:
        def dfs(node: TrieNode, i: int) -> bool:
            if i == len(word):
                return node.is_word  # consumed the whole pattern → must end a word
            ch = word[i]
            if ch == ".":
                # Wildcard: try EVERY child; succeed if any subtree matches the rest.
                for child in node.children.values():
                    if dfs(child, i + 1):
                        return True
                return False
            # Concrete char: follow the single matching edge if it exists.
            nxt = node.children.get(ch)
            return nxt is not None and dfs(nxt, i + 1)

        return dfs(self.root, 0)

When to reach for a Trie

• autocomplete / typeahead / "words with prefix"
• "add and search word" with '.' wildcards (LeetCode 211)
• Word Search II (LeetCode 212): trie of the dictionary, then DFS the board once, pruning whole branches the moment no word shares the current path
• longest common prefix, replace words, maximum XOR (a binary trie)

A quick sanity check on the distinction the is_word flag buys you: after inserting "apple", "app", "application", "banana", search("app") is True (inserted as a full word) but search("appl") is False (a prefix, never a full word) while starts_with("appl") is True.