Word search with a trie

The classic problem is Word Search II (LeetCode 212): given a grid of letters and a list of words, return every word that can be formed by moving between horizontally or vertically adjacent cells, without using the same cell twice in one word.

Searching each word separately repeats the same work: 'oath', 'oat' and 'oats' all explore the same o-a-t paths. A trie merges the words, so one DFS from each cell explores the grid and the trie together. If the letters walked so far are not a prefix of any word, the trie has no child for the next letter and the DFS backs off immediately. That pruning is what makes the problem fast enough.

1. 1Insert all words into a trie. At the end node of each word, store the word itself so you can record it without rebuilding the string.
2. 2For every cell (r, c), start a DFS with the root as the parent node.
3. 3In the DFS, look up the cell's letter among the parent's children. If there is no child, return: no word continues this way.
4. 4If the child node holds a word, add it to the results and remove it from the node so it is not reported twice.
5. 5Mark the cell as visited (for example, replace the letter with '#'), recurse into the four neighbours with the child node, then restore the letter. This is the backtracking step.
6. 6After returning, if the child node has no children and no word left, delete it from its parent. Later searches then stop even earlier, and once all words are found the trie shrinks to nothing.

def find_words(board, words):
    root = {}
    for w in words:                        # build the trie; '$' holds a full word
        node = root
        for ch in w:
            node = node.setdefault(ch, {})
        node['$'] = w
    rows, cols, found = len(board), len(board[0]), []

    def dfs(r, c, parent):
        ch = board[r][c]
        node = parent.get(ch)
        if node is None:
            return                         # no word continues this way: prune
        if '$' in node:
            found.append(node.pop('$'))    # record once, never again
        board[r][c] = '#'                  # mark visited for this path
        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nr, nc = r + dr, c + dc
            if 0 <= nr < rows and 0 <= nc < cols and board[nr][nc] != '#':
                dfs(nr, nc, node)
        board[r][c] = ch                   # undo the mark (backtrack)
        if not node:
            parent.pop(ch)                 # branch has no words left: trim it

    for r in range(rows):
        for c in range(cols):
            dfs(r, c, root)
    return found

The trie is nested dicts with '$' holding a finished word. dfs carries the parent trie node, prunes when the board letter is not a child, records and removes words once found, marks the cell with '#' while exploring and restores it afterwards, and trims trie branches that become empty.

Cost: Worst case O(R x C x 4 x 3^(L-1)) for an R x C board and maximum word length L (4 directions for the first step, at most 3 afterwards because you cannot go back). Building the trie is O(total characters in the words). Extra space O(total characters) for the trie plus O(L) recursion depth.