Code Implementation

class Solution:
    def wordSquares(self, words):
        from collections import defaultdict

        N = len(words[0])
        prefix_map = defaultdict(list)
        for word in words:
            for i in range(N+1):
                prefix_map[word[:i]].append(word)

        results = []

        def backtrack(square):
            if len(square) == N:
                results.append(square[:])
                return
            prefix = ''.join([word[len(square)] for word in square])
            for candidate in prefix_map.get(prefix, []):
                square.append(candidate)
                backtrack(square)
                square.pop()

        for word in words:
            backtrack([word])
        return results
      

class Solution {
public:
    vector<vector<string>> wordSquares(vector<string>& words) {
        int N = words[0].size();
        unordered_map<string, vector<string>> prefix_map;
        for (const string& word : words) {
            for (int i = 0; i <= N; ++i) {
                prefix_map[word.substr(0, i)].push_back(word);
            }
        }

        vector<vector<string>> results;
        vector<string> square;

        function<void()> backtrack = [&]() {
            if (square.size() == N) {
                results.push_back(square);
                return;
            }
            string prefix;
            for (int i = 0; i < square.size(); ++i) {
                prefix += square[i][square.size()];
            }
            for (const string& candidate : prefix_map[prefix]) {
                square.push_back(candidate);
                backtrack();
                square.pop_back();
            }
        };

        for (const string& word : words) {
            square = {word};
            backtrack();
        }
        return results;
    }
};
      

class Solution {
    public List<List<String>> wordSquares(String[] words) {
        int N = words[0].length();
        Map<String, List<String>> prefixMap = new HashMap<>();
        for (String word : words) {
            for (int i = 0; i <= N; ++i) {
                String prefix = word.substring(0, i);
                prefixMap.computeIfAbsent(prefix, k -> new ArrayList<>()).add(word);
            }
        }

        List<List<String>> results = new ArrayList<>();
        List<String> square = new ArrayList<>();

        backtrack(words, N, prefixMap, results, square);
        return results;
    }

    private void backtrack(String[] words, int N, Map<String, List<String>> prefixMap,
                           List<List<String>> results, List<String> square) {
        if (square.size() == N) {
            results.add(new ArrayList<>(square));
            return;
        }
        StringBuilder prefix = new StringBuilder();
        int idx = square.size();
        for (String s : square) {
            prefix.append(s.charAt(idx));
        }
        for (String candidate : prefixMap.getOrDefault(prefix.toString(), new ArrayList<>())) {
            square.add(candidate);
            backtrack(words, N, prefixMap, results, square);
            square.remove(square.size() - 1);
        }
    }
}
      

var wordSquares = function(words) {
    const N = words[0].length;
    const prefixMap = {};
    for (let word of words) {
        for (let i = 0; i <= N; ++i) {
            let prefix = word.substring(0, i);
            if (!prefixMap[prefix]) prefixMap[prefix] = [];
            prefixMap[prefix].push(word);
        }
    }

    const results = [];

    function backtrack(square) {
        if (square.length === N) {
            results.push(square.slice());
            return;
        }
        let prefix = '';
        let idx = square.length;
        for (let i = 0; i < idx; ++i) {
            prefix += square[i][idx];
        }
        let candidates = prefixMap[prefix] || [];
        for (let candidate of candidates) {
            square.push(candidate);
            backtrack(square);
            square.pop();
        }
    }

    for (let word of words) {
        backtrack([word]);
    }
    return results;
};
      

Problem Description

Given an array of strings words, all of the same length, you are to find all possible word squares that can be built from them. A word square is a set of words arranged in a square grid so that the words read the same horizontally and vertically. That is, the word at row i is the same as the word at column i.

Constraints:

• You may use each word from words as many times as you like.
• All words have the same length N. The square is N x N.
• Return all possible valid word squares using the given words.

Example:
Input: ["area","lead","wall","lady","ball"]
Output: [ ["wall","area","lead","lady"], ["ball","area","lead","lady"] ]

Thought Process

At first glance, the problem seems to require generating all possible arrangements of words in a square, then checking if the vertical and horizontal words match. This brute-force approach would involve trying every possible combination, which can be very slow as the number of words and their length grows.

However, we can notice that for a word square, the k-th word must have its first k letters match the k-th column so far. This means that at each step, we can use the current partial square to determine a prefix and only consider words that match this prefix for the next row. This insight allows us to prune invalid combinations early.

To do this efficiently, we can build a prefix map (like a dictionary or trie) that lets us quickly find all words starting with a given prefix. Using backtracking, we can build the square row by row, always ensuring the columns match.

Solution Approach

• Step 1: Build a Prefix Map
  • Create a map (dictionary) where each prefix maps to a list of words starting with that prefix.
  • For every word, for every possible prefix (including the empty prefix), add the word to the prefix's list.
• Step 2: Backtracking Function
  • Start with each word as the first row of the square.
  • At each level of recursion, build the prefix for the next row by collecting the k-th letter of each word in the current square (where k is the current depth).
  • Use the prefix map to find all words that can be the next row.
  • For each candidate, add it to the current square and recurse. If the square is complete (length N), add it to the results.
  • Backtrack by removing the last word and trying the next candidate.
• Step 3: Collect Results
  • Return all completed squares found by the backtracking process.

Using a prefix map allows us to efficiently find only those words that can possibly fit, drastically reducing the number of combinations we need to check.

Example Walkthrough

Let's use the sample input ["area","lead","wall","lady","ball"] (all words of length 4).

1. Start with "wall":
   Square so far: ["wall"]
   • To fill row 2, collect the prefix for column 2: "a" (from "wall"[1]).
   • Possible words starting with "a": ["area"]
2. Add "area":
   Square so far: ["wall", "area"]
   • To fill row 3, collect the prefix for column 3: "le" (from "wall"[2] and "area"[2]).
   • Possible words starting with "le": ["lead"]
3. Add "lead":
   Square so far: ["wall", "area", "lead"]
   • To fill row 4, collect the prefix for column 4: "lad" (from "wall"[3], "area"[3], "lead"[3]).
   • Possible words starting with "lad": ["lady"]
4. Add "lady":
   Square so far: ["wall", "area", "lead", "lady"]
   • Square is complete and valid (columns and rows match).
5. Backtrack and try "ball" as the first word:
   • Repeat the process, and you find that ["ball", "area", "lead", "lady"] is also a valid square.
6. Continue for all words as the first row.

This process efficiently finds all valid word squares.

Time and Space Complexity

Brute-force approach: Try all possible combinations of N words for the square, check if columns match. If there are W words of length N, the number of combinations is W^N. Checking validity for each square takes O(N^2) time.

Optimized approach (with prefix map):

• Building the prefix map: O(W * N^2) (since each word has N prefixes, each of length up to N).
• Backtracking: The number of recursive calls is much less than W^N because at each step, only words matching the current prefix are considered. In the worst case (if all prefixes are present), it's still exponential, but in practice, the prefix map prunes the search space significantly.
• Space: The prefix map uses O(W * N^2) space. The recursion stack uses up to O(N) space.

Summary

The word squares problem is a classic example of using backtracking with efficient pruning. By recognizing that each new row must match the current columns, and using a prefix map to quickly find candidates, we avoid brute-force enumeration of all possibilities. This makes the solution elegant and efficient, and highlights the power of combining recursion with smart data structures.