LC 0425 [H] Word Squares - ALawliet/algorithms GitHub Wiki

class Solution(object):
    def wordSquares(self, words):
        """
        :type words: List[str]
        :rtype: List[List[str]]
        """
        def search(square):
            ROWS, COLS = len(square), len(square[0])
            
            if ROWS == COLS:    # solution found
                ans.append(square)
            else:
                req_prefix = ""    # Generate the required prefix
                for r in range(ROWS):
                    req_prefix += square[r][ROWS]
                    
                if req_prefix in prefix_to_words:    # Perform backtracking
                    for word in prefix_to_words[req_prefix]:
                        search(square + [word])
            
        # Generate all possible prefixes
        prefix_to_words = defaultdict(list) # prefix -> words
        for word in words:
            for i in range(len(word)):
                prefix = word[:i]
                prefix_to_words[prefix].append(word)
        
        ans = []
        for word in words:
            search([word])
        return ans