# 用Hash or Trie树记录下以某个前缀开头的单词有哪些
# 比如以l开头的有lead, lady ,以le开头的有lead, 以lea开头的有lead

# 第一个单词填了ball，第二个词想填area的话，字典中必须有以le,la开头的单词，否则如果没有的话，不能填area
# b a l l
# a r e a
# l e a d
# l a d y

# l --> lead, lady
# le -> lead
# lea -> lead
# lead -> lead

# 用Hash or Trie树记录下以某个前缀开头的单词有哪些
# 比如以l开头的有lead, lady ,以le开头的有lead, 以lea开头的有lead

class TrieNode:
    def __init__(self):
        self.children = {}
        self.word_list = []

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

    def add(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_list.append(word)

    def get_words_with_prefix(self, prefix):
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return []
            node = node.children[ch]
        return node.word_list


class Solution:
    """
    @param: words: a set of words without duplicates
    @return: all word squares
    """

    def wordSquares(self, words):
        trie = Trie()
        for word in words:
            trie.add(word)

        squares = []
        for word in words:
            # not trie.root , trie.root is TrieNode
            self.search(trie, [word], squares)

        return squares

    def search(self, trie, square, squares):
        # n = len(square[0])
        # curt_index = len(square)
        # if curt_index == n:
        #     squares.append(list(square))
        #     return

        # exit
        # if solution square 的length 等于单词的长度，那么就加这个solution到最后的答案里面
        if len(square) == len(square[0]):
            squares.append(list(square))
            return


        # b a l l
        # a r e a
        # 在预测第三个加入的单词的时候，需要从前缀"le"的单词里找。怎么得到prefix的呢？
        # 注意到：加第一个单词的时候，word的index0 postion的字母已经找到了，下一个search应该从index1开始
        # 所以我们如果想加第N个单词的时候，下一个字母应该是indexN.
        # 我们已经找到了BALL, AREA,当我们需要预测第三个单词的时候， 我们应该取BALL的INDEX3的字母（L） + Area的
        # index3的字母E， 这两个字母组成"le"前缀。

        #find prefix for next search, 下一个index 从len(square)开始
        # ball找好以后，下一个应该是以a开头的单词
        # ball和area填好以后，下一个应该是le开头的单词
        tmp_list = []
        next_index = len(square)
        for i in range(next_index):
            tmp_list.append(square[i][next_index])
        prefix = ''.join(tmp_list)

        # search
        for word in trie.get_words_with_prefix(prefix):
            square.append(word)
            self.search(trie, square, squares)
            square.pop()  # remove the last word


sol = Solution()
words = ["ball","lead","wall","lady","area"]
res = sol.wordSquares(words)
print(res)