Leetcode: 212. Word Search II

Mistake: High constant factor in the implementation. Did not implement any prune to the backtracking.
Can we pre-process the input in a way to make easy to solve the problem? The problem requires backtracking the board to find words. Instead of look up the words on a list, we create a Trie with all words on the dictionary and backtrack while we traverse the tree. Words found are removed from the tree to avoid finding the same word over and over again. Time complexity is \(O(4 \times 10 \times n^2)\), and space complexity is \(O(\sum|word[i]|)\).

from typing import List


class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        N = len(board)
        M = len(board[0])

        root = {}
        for w in words:
            node = root
            for c in w:
                node.setdefault(c, {})
                node = node[c]
            node["$"] = w

        ans = set()

        def dfs(i, j, node):
            if node.get("$"):
                ans.add(node["$"])
                del node["$"]

            c = board[i][j]
            board[i][j] = "*"
            for di, dj in [[+0, +1], [+0, -1], [+1, +0], [-1, +0]]:
                ni = di + i
                nj = dj + j
                if 0 <= ni < N and 0 <= nj < M and board[ni][nj] in node:
                    dfs(ni, nj, node[board[ni][nj]])
                    if len(node[board[ni][nj]]) == 0:
                        del node[board[ni][nj]]
            board[i][j] = c

        for i in range(N):
            for j in range(M):
                if board[i][j] in root:
                    dfs(i, j, root[board[i][j]])

        return list(ans)


assert Solution().findWords(
    [
        ["o", "a", "a", "n"],
        ["e", "t", "a", "e"],
        ["i", "h", "k", "r"],
        ["i", "f", "l", "v"],
    ],
    ["oath", "pea", "eat", "rain"],
) == ["eat", "oath"]
assert Solution().findWords([["a", "b"], ["c", "d"]], ["abcb"]) == []