Word Search II: Trie + Backtracking on a Grid

Word search, the single-word version, is a clean backtracking exercise: start at every cell, walk to neighbours, and undo each step on the way back. Word search II hands you a whole list of words and the same board, and the obvious answer, running word search once per word, repeats nearly all of its work. The fix is to stop searching for words and start searching for prefixes. Put the list in a trie, and one walk over the grid advances every candidate word at the same time.

The problem it solves

Given a grid of letters and a list of words, return every word that can be traced through horizontally or vertically adjacent cells, using each cell at most once per word.

Running single-word search for each of w words costs w full backtracking searches. Two words that share a prefix, such as sea and seat, are traced along exactly the same cells twice. With hundreds of words over a small alphabet, most of the effort is repeated walks down the same paths.

The intuition

A trie stores words as paths from the root, one letter per edge, with shared prefixes stored once. That makes it the right partner for a grid walk, because the walk also builds a word one letter at a time.

So the recursion carries a trie node instead of an index into one word. At each cell it asks one question: is this letter a child of the current node? If not, no word in the list starts with the letters on this path, and the whole branch is dropped on the spot. If it is, step into that child, record a word if the child marks the end of one, block the cell, try the four neighbours, and unblock the cell on the way back.

Two small refinements make the full solution fast in practice:

• Remove a word once it is found. Popping the end marker means the same word is never reported twice and never searched for again.
• Prune exhausted branches. When a trie node has no children left, delete it from its parent, so later walks die one step earlier.

The trie is one of the patterns on the patterns cheat sheet, and this board search is its standard showcase.

Watch it run

The animation puts two words in a trie, car and cat, over a two-row grid, and every cell is a possible first letter. The walk starts at c; the root has a c child, so this walk is alive. It steps to a, still a prefix, with both words in play. Then r is a child of a and carries the end flag, which gives car. The walk backs up to a and takes the other child instead, which gives cat. Then it starts at x: the root has no x child, so the walk dies on its first comparison. The same happens for y, and neither cell's neighbours are ever visited, so that whole branch is gone. The final frame sums it up: one walk over the grid found both words, and every dead start cost a single lookup.

The code

The lesson's version, with the trie written out by hand as nested dictionaries and $ as the end-of-word flag:

grid = [["c", "a", "r"], ["x", "t", "y"]]
trie = {"c": {"a": {"r": {"$": 1}, "t": {"$": 1}}}}
found = set()
def dfs(r, c, node, word):
    ch = grid[r][c]
    if ch not in node: return          # dead prefix — prune this branch
    node, word = node[ch], word + ch
    if "$" in node: found.add(word)
    grid[r][c] = "#"                   # no cell twice inside one word
    for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
        if 0 <= r + dr < 2 and 0 <= c + dc < 3:
            dfs(r + dr, c + dc, node, word)
    grid[r][c] = ch                    # put it back for other paths
for r, c in [(r, c) for r in range(2) for c in range(3)]: dfs(r, c, trie, "")
print(sorted(found))                   # ['car', 'cat']

The general solution builds the trie from the word list, stores each word at its end node, removes it once found, and deletes branches that have nothing left. A counter records how many cells the walk visits. The second example shows the "each cell once" rule: aba would need the a twice:

def build_trie(words):
    root = {}
    for w in words:
        node = root
        for ch in w:
            node = node.setdefault(ch, {})
        node["$"] = w                  # store the word itself at its end
    return root

calls = {"trie": 0, "one word": 0}

def find_words(board, words):
    rows, cols = len(board), len(board[0])
    root, out = build_trie(words), []
    def walk(r, c, parent):
        calls["trie"] += 1
        ch = board[r][c]
        node = parent.get(ch)
        if node is None:
            return                     # not a prefix of any remaining word
        if "$" in node:
            out.append(node.pop("$"))  # report once, then forget it
        board[r][c] = "#"
        for nr, nc in ((r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)):
            if 0 <= nr < rows and 0 <= nc < cols and board[nr][nc] != "#":
                walk(nr, nc, node)
        board[r][c] = ch
        if not node:
            del parent[ch]             # branch exhausted: prune it for good
    for r in range(rows):
        for c in range(cols):
            walk(r, c, root)
    return out

board = [list("seat"), list("tonr"), list("ards")]
print(sorted(find_words(board, ["sea", "seat", "toe", "tan", "and", "note", "rats", "stone"])))
# ['and', 'rats', 'sea', 'seat', 'tan', 'toe']
print(find_words([list("ab")], ["aba", "ab", "ba"]))               # ['ab', 'ba']

For comparison, single-word search, the plain backtracking version, run once per word. On a random 6 by 6 board with 457 distinct five-letter words, both find the same 123 words, and the trie walk visits about sixteen times fewer cells:

def exists(board, word):               # word search I: one word, plain backtracking
    rows, cols = len(board), len(board[0])
    def go(r, c, i):
        calls["one word"] += 1
        if board[r][c] != word[i]:
            return False
        if i == len(word) - 1:
            return True
        ch, board[r][c] = board[r][c], "#"
        hit = any(go(nr, nc, i + 1)
                  for nr, nc in ((r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1))
                  if 0 <= nr < rows and 0 <= nc < cols)
        board[r][c] = ch
        return hit
    return any(go(r, c, 0) for r in range(rows) for c in range(cols))

import random

random.seed(5)
board = [[random.choice("abcde") for _ in range(6)] for _ in range(6)]
words = sorted({"".join(random.choice("abcde") for _ in range(5)) for _ in range(500)})
calls = {"trie": 0, "one word": 0}
hits = find_words([row[:] for row in board], words)
per_word = [w for w in words if exists(board, w)]
print(len(words), len(hits), len(per_word), calls)
# 457 123 123 {'trie': 2099, 'one word': 33601}

Then the trie search checked against one-word search on 1,500 random small boards and word lists, including one-row and one-column boards, with no word reported twice:

random.seed(22)
ok = True
for _ in range(1500):
    rows, cols = random.randint(1, 4), random.randint(1, 4)
    board = [[random.choice("abc") for _ in range(cols)] for _ in range(rows)]
    words = sorted({"".join(random.choice("abc") for _ in range(random.randint(1, 5)))
                  for _ in range(random.randint(1, 8))})
    want = sorted(w for w in words if exists(board, w))
    got = find_words([row[:] for row in board], words)
    ok &= sorted(got) == want and len(got) == len(set(got))
print(ok)                              # True

The complexity

• Building the trie: O(L), where L is the total number of letters in the word list.
• Searching: from each of the m · n cells, the first step has up to four choices and every later step up to three, since the walk never goes back to the cell it came from. With maximum word length k, the worst case is O(m · n · 4 · 3^(k-1)), and it does not grow with the number of words. In practice pruning keeps it far below that bound.
• Space: O(L) for the trie plus O(k) for the recursion.

Where it goes wrong

• Searching for each word separately. Correct, but it pays for every shared prefix again. That is the whole reason for the trie.
• Forgetting to restore the cell. A blocked cell left behind makes later paths miss words that need it.
• Reporting duplicates. A word that can be traced two ways is found twice unless its end marker is removed, or results are kept in a set.
• Carrying the whole prefix string into every call. The trie node already knows where the walk is; storing the word at its end node avoids building strings on every step.

When it shows up in interviews

It is a standard hard question and the usual follow-up to word search with backtracking. The interviewer wants to hear why a trie beats repeated searches, and whether you restore the board and avoid duplicates. Building the trie itself is the subject of implement a trie, and the trade-off against a hash set of prefixes is in trie vs hash set for prefixes.

How to say it in an interview

"I put all the words into a trie and store each word at its end node. Then I start a depth-first walk from every cell, carrying a trie node instead of an index into one word. If the cell's letter is not a child of the current node, no word starts with this path and I return immediately. Otherwise I step into the child, record a word if one ends there, block the cell, try the four neighbours, and restore it. I remove words once found so they are not reported twice, and I delete trie branches that become empty. The cost depends on the board and the longest word, not on how many words there are."