212. Word Search II - cocoder39/coco39_LC GitHub Wiki

212. Word Search II

backtracking + trie for pruning

class TrieNode:
    def __init__(self):
        self.isWord = False
        self.childres = collections.defaultdict(TrieNode)

class Trie:
    def __init__(self):
        self.root = TrieNode()
    
    def insert(self, word):
        cur = self.root
        for ch in word:
            cur = cur.childres[ch]
        cur.isWord = True

class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        trie = Trie()
        for word in words:
            trie.insert(word)
        
        DIR = [(-1, 0), (1, 0), (0, -1), (0, 1)]
        m, n = len(board), len(board[0])
        
        def dfs(i, j, word, cur):
            ch = board[i][j]
            if ch not in cur.childres:
                return
            
            visited.add((i, j))
            if cur.childres[ch].isWord:
                res.append(''.join(word))
                cur.childres[ch].isWord = False # avoid duplication
            
            for x, y in DIR:
                if 0<=i+x<m and 0<=j+y<n and (i+x, j+y) not in visited:
                    word.append(board[i+x][j+y])
                    dfs(i+x, j+y, word, cur.childres[ch])
                    word.pop()
            visited.remove((i, j))
                    
        visited = set()
        res = []
        for i in range(m):
            for j in range(n):
                dfs(i, j, [board[i][j]], trie.root)
        return res