Problem statement

https://leetcode.com/problems/word-squares/

Solution

One possible way to solve the problem is to generate all possible prefixes. Then we use dfs with words, k parameters, where ther first one is the words we built so far and the second is number of words we have. Each time when we try to add new word we need to look at k symbol of each generated word and try to find word starting from this prefix.

Complexity

Code

class Solution:
    def wordSquares(self, words):
        n, m = len(words), len(words[0])
        prefs, out = defaultdict(set), []

        for word in words:
            for i in range(m):
                prefs[word[:i]].add(word)
        
        def dfs(words, k):
            if k == m:
                out.append(words)
                return
                
            pref = "".join(words[i][k] for i in range(k))   
            for cand in prefs[pref]: 
                dfs(words + [cand], k+1)
                
        dfs([], 0)
        return out