'''
主要思路： 如果当前走的路径不是前缀树的前缀，我们就可以提前结束了。

从图中的每个位置出发 -> 看遍历过程中是否遇到了 words 中的某个单词
遍历过程中判断是否遇到了某个单词，我们可以事先把所有单词存到前缀树中。
这样的话，如果当前走的路径不是前缀树的前缀，我们就可以提前结束了。如果是前缀树的中的单词，我们就将其存到结果中。

至于实现的话，我们可以在遍历过程中，将当前路径的单词传进函数，然后判断当前路径构成的单词是否是在前缀树中出现即可。
这个想法可行，但不够好，因为每次都从前缀树中判断当前路径的单词，会带来重复的判断。比如先判断了 an 存在于前缀树中，接下来假如路径变成 ang ，
判断它在不在前缀中，又需要判断一遍 an 。

因此，我们可以将前缀树融合到我们的算法中，递归中去传递前缀树的节点，判断当前节点的孩子是否为 null，如果是 null 说明当前前缀不存在，可以提前结束。
如果不是 null，再判断当前节点是否是单词的结尾，如果是结尾直接将当前单词加入。

由于递归过程中没有加路径，所以我们改造一下前缀树的节点，将单词直接存入节点，这样的话就可以直接取到了。
'''
class TrieNode:
    # key : value ==> character : TrieNode
    def __init__(self):
        self.children = {}
        self.word = None

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def addNode(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.word = word # save word

class Solution:
    directions = [(0,1), (0, -1), (1, 0), (-1, 0)]

    def wordSearchII(self, board, words):
        if not words:
            return []

        # build Trie tree for words
        tree = Trie()
        for word in words:
            tree.addNode(word)

        # DFS search in board and trie tree
        # if result =[], then  ['dog', 'dad', 'again', 'dad', 'can'], has duplicates
        result = set()
        # result = []
        self.m, self.n = len(board), len(board[0])
        for i in range(self.m):
            for j in range(self.n):
                self.search(board, i, j, tree.root, result)
        return list(result) # convert it to list
        # return result

    def search(self, board, i, j, node, result):
        dirs = [(0, 1), (0, -1), (1, 0), (-1, 0)]
        if i < 0 or i >= len(board) or j < 0 or j >= len(board[0]):
            return

        ch = board[i][j]
        if ch not in node.children:
            return

        # 找到，加入结果中
        if node.children[ch].word is not None:
            result.add(node.children[ch].word)
            # result.append(node.children[ch].word)

        # avoid to duplicated search
        board[i][j] = '0'
        for k in range(4):
            x, y = i + dirs[k][0], j + dirs[k][1]
            self.search(board, x, y, node.children[ch], result)
        # revert the mark for next search
        board[i][j] = ch

    def search1(self, board, x, y, node, result):
        # 越界 或者已经遍历过，return
        if x < 0 or x >= self.m or y < 0 or y >= self.n or board[x][y] == '0':
            return

        # 没有找到，return
        ch = board[x][y]
        if ch not in node.children:
            return

        # 找到，加入结果中
        if node.children[ch].word is not None:
            result.add(node.children[ch].word)

        # 从board[x][y]这个点开始，四个方向遍历
        for i, j in self.directions:
            board[x][y] = '0'  # avoid to duplicated search
            self.search(board, x + i, y + j, node.children[ch], result)
            board[x][y] = ch # revert the mark for next search

    # search 也可以写成一下形式，在循环里判断越界
    def search2(self, board, x, y, node, result):
        # 没有找到，return
        ch = board[x][y]
        if ch not in node.children:
            return

        # 找到，加入结果中
        if node.children[ch].word is not None:
            result.add(node.children[ch].word)

        # 从board[x][y]这个点开始，四个方向遍历
        for i, j in self.directions:
            if x + i < 0 or x + i >= self.m or y + j < 0 or y + j >= self.n or board[x + i][y + j] == '0':
                continue # use continue NOT return
            board[x][y] = '0'  # avoid to duplicated search
            self.search(board, x + i, y + j, node.children[ch], result)
            board[x][y] = ch # revert the mark for next search

if __name__ == "__main__":
    sol = Solution()
    board = [['d','o','a','f'],['a','g','a','i'],['d','c','a','n']]
    words = ["dog","dad","dgdg","can","again"]
    res = sol.wordSearchII(board, words)
    print(res)