Word Search II

class TrieNode: def __init__(self): self.children = {} self.end = False self.word = None class Trie: def __init__(self): self.root = TrieNode() def add_word(self, word): cur = self.root for ch in word: if ch not in cur.children: cur.children[ch] = TrieNode() cur = cur.children[ch] cur.end = True class Solution: def findWords(self, board: List[List[str]], words: List[str]) -> List[str]: nrows = len(board) ncols = len(board[0]) ans = set() # Build trie trie = Trie() for word in words: trie.add_word(word) # Define DFS function vis = set() def dfs(r, c, cur, word): # if r < 0 or c < 0 or r == nrows or c == ncols or board[r][c] not in cur.children or (r, c) in vis: # return if r >= 0 and r < nrows and c >= 0 and c < ncols and board[r][c] in cur.children and (r, c) not in vis: vis.add((r, c)) next_cur = cur.children[board[r][c]] word += board[r][c] if next_cur.end: ans.add(word) for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]: nr, nc = r + dr, c + dc dfs(nr, nc, next_cur, word) vis.remove((r, c)) # Conduct DFS for r in range(nrows): for c in range(ncols): dfs(r, c, trie.root, '') return list(ans)