425. Word Squares

利用Trie和backtracking来做。利用trie来找prefix

class Solution(object):
    def wordSquares(self, words):
        """
        :type words: List[str]
        :rtype: List[List[str]]
        """
        # Write your code here
        if not words:
            return []
        root = self.buildTree(words)
        self.res = []
        self.dfs(root, [], words)
        return self.res
        
    
    def dfs(self, root, cur, words):
        if len(cur) == len(words[0]):
            self.res.append(["".join(w) for w in cur])
            return True
        
        prefix = self.getPrefix(cur)
        
        for word in self.getPrefixWords(root, prefix):
            cur.append(list(word))
            self.dfs(root, cur, words)
            cur.pop()
    
    def getPrefix(self, cur):
        pos = len(cur)
        prefix = ""
        for row in cur:
            prefix += row[pos]
        return prefix
        
    def getPrefixWords(self, root, prefix):
        # get all words with current prefix
        node = root
        for c in prefix:
            if c not in node.children:
                return []
            node = node.children[c]
        # start from current node, get all node.stop True
        return self.getWords(node, prefix, "", [])
    
    def getWords(self, node, prefix, cur, res):
        if node.stop:
            res.append(prefix+cur)
        
        for n in node.children:
            self.getWords(node.children[n], prefix, cur+n, res)
        return res

    
    def buildTree(self, words):
        root = TrieNode("")
        for word in words:
            node = root
            for c in word:
                if c not in node.children:
                    node.children[c] = TrieNode(c)
                node = node.children[c]
            node.stop = True
        
        return root


class TrieNode(object):
    def __init__(self, c):
        self.c = c
        self.stop = False
        self.children = {}