[leetcode]Word Squares

class Solution:
    def wordSquares(self, words: List[str]) -> List[List[str]]:
        wordsDict = {}
        
        for i in range(len(words)):
            word = words[i]
            for k in range(len(word)):
                if word[0:k+1] not in wordsDict:
                    wordsDict[word[0:k+1]] = []
                wordsDict[word[0:k+1]].append(word)
            
        result = []
        
        # backtrace
        for i in range(len(words)):
            self.backtrace([words[i]], result, words, wordsDict)
            
        return result
        
    def backtrace(self, square, result, words, wordsDict):
        n = len(words[0])
        i = len(square) - 1 # idx for newly add line
        for k in range(0, i):
            if square[i][k] != square[k][i]:
                # check fail, return
                return
        
        if i + 1 == n: # square is good
            result.append(square.copy())
            return
        
        prefix = ‘‘
        for k in range(0, i+1):
            prefix += square[k][i+1]

        if prefix not in wordsDict:
            return
        
        for word in wordsDict[prefix]:
            square.append(word)
            self.backtrace(square, result, words, wordsDict)
            square.pop()