Solution Explanation for LeetCode 425: Word Squares

This problem asks us to find all possible word squares that can be formed using a given list of unique words. A word square is a set of words where the kth row and column read the same string. The solution uses a Trie data structure and backtracking.

Approach

1. Trie Construction: A Trie (prefix tree) is built from the input words. Each node in the Trie represents a prefix, and each node stores a list of indices of words in the input list that share that prefix. This allows efficient prefix searching.

2. Backtracking (DFS): A depth-first search (DFS) algorithm is used to explore possible word square combinations. At each level of the DFS, it attempts to add a word to the current partial word square.

3. Prefix Check: Before adding a word, the algorithm checks if a valid prefix exists for the next position in the word square. The prefix is formed by concatenating the characters at the current depth from all the already selected words. This prefix is then searched in the Trie to find words that start with that prefix.

4. Base Case: The base case of the DFS is reached when a complete word square of size n x n (where n is the length of each word) is constructed. This complete word square is added to the result list.

Code Explanation (Python)

class Trie:
    def __init__(self):
        self.children = [None] * 26
        self.v = []  # Indices of words with this prefix
 
    def insert(self, w, i):
        node = self
        for c in w:
            idx = ord(c) - ord('a')
            if node.children[idx] is None:
                node.children[idx] = Trie()
            node = node.children[idx]
            node.v.append(i)
 
    def search(self, w):
        node = self
        for c in w:
            idx = ord(c) - ord('a')
            if node.children[idx] is None:
                return []
            node = node.children[idx]
        return node.v
 
 
class Solution:
    def wordSquares(self, words: List[str]) -> List[List[str]]:
        def dfs(t):  # t is the current partial word square
            if len(t) == len(words[0]):  # Base case: complete word square
                ans.append(t[:])
                return
            idx = len(t)
            pref = [v[idx] for v in t]  # Prefix for the next word
            indexes = trie.search(''.join(pref))
            for i in indexes:
                t.append(words[i])  # Add the word
                dfs(t)  # Explore further
                t.pop()  # Backtrack
 
        trie = Trie()
        ans = []
        for i, w in enumerate(words):
            trie.insert(w, i)
        for w in words:  # Start DFS with each word
            dfs([w])
        return ans
 

The Java and Go code follow a very similar structure and logic.

Time and Space Complexity Analysis

• Time Complexity: O(M * N2 * 4N), where M is the number of words and N is the length of each word. Building the Trie takes O(M * N). The DFS explores a branching factor of up to 4N (each position in the word square can potentially have 4 choices). For each node explored in the DFS, searching in the Trie takes O(N) time, and adding/removing a word takes O(N).

• Space Complexity: O(M * N) for the Trie and O(N2) for the space used by the DFS call stack and result list. In the worst case, the Trie could potentially store all prefixes of all words. The call stack depth is bounded by N, and the result list stores at most 4N word squares. Therefore, the space complexity is dominated by the Trie.

In essence, this solution cleverly uses a Trie for efficient prefix searching and backtracking to explore all possible word square combinations. The time complexity is exponential due to the nature of the problem, but the Trie helps to improve search efficiency compared to naive approaches.