Wyszukiwanie słów II

class TrieNode:
    def __init__(self):
        self.children = {}
        self.isWord = False
    def addWord(self, word):
        curr = self
        for c in word:
            if c not in curr.children: curr.children[c] = TrieNode()
            curr = curr.children[c]
        curr.isWord = True

def find_words_opt(board, words):
    root = TrieNode()
    for w in words: root.addWord(w)
    ROWS, COLS = len(board), len(board[0])
    res, visit = set(), set()
    def dfs(r, c, node, word):
        if r < 0 or r == ROWS or c < 0 or c == COLS or (r, c) in visit or board[r][c] not in node.children:
            return
        visit.add((r, c))
        node = node.children[board[r][c]]
        word += board[r][c]
        if node.isWord: res.add(word)
        for dr, dc in [[0, 1], [0, -1], [1, 0], [-1, 0]]:
            dfs(r + dr, c + dc, node, word)
        visit.remove((r, c))
    for r in range(ROWS):
        for c in range(COLS):
            dfs(r, c, root, "")
    return list(res)

def find_words_brute(board, words):
    def exist(word):
        ROWS, COLS = len(board), len(board[0])
        def dfs(r, c, i, visited):
            if i == len(word): return True
            if r < 0 or c < 0 or r >= ROWS or c >= COLS or (r, c) in visited or board[r][c] != word[i]:
                return False
            visited.add((r, c))
            res = any(dfs(r+dr, c+dc, i+1, visited) for dr, dc in [[1,0],[-1,0],[0,1],[0,-1]])
            visited.remove((r, c))
            return res
        return any(dfs(r, c, 0, set()) for r in range(ROWS) for c in range(COLS))
    return [w for w in words if exist(w)]