#Using dfs to search for all possible solutions
#given the row from 0 to i-1, we know the prefix of the word to fill in the row i
#We can then search for the word with this prefix
#In order to search it, we use trie.
#The difference between trie used here and common trie API is that a Prefix List is required for each node
#The Prefix List stores all the words index which has the certain prefix

#Let n = len(words[0]) m = len(words)
#TC: For constructing trie, it needs O(m*26*n)
#For backtracking: maximum O(m*n)
#SC: O(m*26*n)

class Node:
    def __init__(self):
        self.PrefixIdx = []
        self.Dict = {}
        self.isEnd = False

class Trie:
    def __init__(self):
        self.dummy = Node()
    def insert(self, word, idx):
        head = self.dummy
        for c in word:
            if c not in head.Dict:
                head.Dict[c] = Node()
            head = head.Dict[c]
            head.PrefixIdx.append(idx)

        head.isEnd = True

class Solution:
    def wordSquares(self, words: List[str]) -> List[List[str]]:
        trie = Trie()
        n = len(words[0])
        for idx, w in enumerate(words):
            trie.insert(w, idx)

        res = []

        def dfs(WordList):

            if len(WordList) == n:
                res.append(WordList)
                return

            TotalRow = len(WordList)
            prefix = ''
            for i in range(TotalRow):
                prefix += WordList[i][TotalRow]

            head = trie.dummy
            for c in prefix:
                if c not in head.Dict:
                    return
                else:
                    head = head.Dict[c]

            for w_idx in head.PrefixIdx:
                WordList.append(words[w_idx])
                dfs(WordList.copy())
                WordList.pop()

        for w in words:
            dfs([w])

        return res