Skip to content

425. Word Squares

Given an array of unique strings words, return all the word squares you can build from words. The same word from words can be used multiple times. You can return the answer in any order.

A sequence of strings forms a valid word square if the kth row and column read the same string, where 0 <= k < max(numRows, numColumns).

  • For example, the word sequence ["ball","area","lead","lady"] forms a word square because each word reads the same both horizontally and vertically.

Example 1:

Input: words = ["area","lead","wall","lady","ball"]
Output: [["ball","area","lead","lady"],["wall","area","lead","lady"]]
Explanation:
The output consists of two word squares. The order of output does not matter (just the order of words in each word square matters).

Example 2:

Input: words = ["abat","baba","atan","atal"]
Output: [["baba","abat","baba","atal"],["baba","abat","baba","atan"]]
Explanation:
The output consists of two word squares. The order of output does not matter (just the order of words in each word square matters).

Constraints:

  • 1 <= words.length <= 1000
  • 1 <= words[i].length <= 4
  • All words[i] have the same length.
  • words[i] consists of only lowercase English letters.
  • All words[i] are unique.

Solution:

class Solution {
    static class TrieNode{
        Map<Character, TrieNode> children = new HashMap<>();
        List<String> wordsWithPrefix = new ArrayList<>(); 

    }

    static class Trie{
        TrieNode root;

        Trie(){
            root = new TrieNode();
        }

        // Insert words into the Trie
        public void insert(String word){
            TrieNode cur = root;
            for (char c : word.toCharArray()){
                cur.wordsWithPrefix.add(word);
                cur.children.putIfAbsent(c, new TrieNode());
                cur = cur.children.get(c);
            }
            cur.wordsWithPrefix.add(word); // At the end of the word, add it to the list

        }

        // Find all words with the given prefix
        public List<String> findByPrefix(String prefix){
            TrieNode node = root;
            for (char c : prefix.toCharArray()){
                if (!node.children.containsKey(c)){
                    return new ArrayList<>(); // Return empty if no words with this prefix
                }
                node = node.children.get(c);
            }

            return node.wordsWithPrefix; // Return all words that have this prefix
        }
    }

    public List<List<String>> wordSquares(String[] words) {
        List<List<String>> result = new ArrayList<>();

        if (words == null || words.length == 0){
            return result;
        }

        Trie trie = new Trie();
        for (String word : words){
            trie.insert(word);
        }

        // Backtrck to build word squares
        List<String> subResult = new ArrayList<>();
        for (String word : words){
            subResult.add(word);
            backtrack(1, subResult, trie, result); // 
            subResult.remove(subResult.size() - 1); 
        }

        return result;


    }

    private void backtrack(int step, List<String> subResult, Trie trie, List<List<String>> result){
        int wordLength = subResult.get(0).length();

        if (step == wordLength){
            result.add(new ArrayList<>(subResult));
            return;
        }

        StringBuilder prefixBuilder = new StringBuilder();

        for (String word : subResult){
            prefixBuilder.append(word.charAt(step));
        }

        String prefix = prefixBuilder.toString();

        List<String> candidates = trie.findByPrefix(prefix);
        for (String candidate : candidates){
            subResult.add(candidate);
            backtrack(step + 1, subResult, trie, result);
            subResult.remove(subResult.size() - 1);
        }
    }
}
graph TD
    Root((Root))

    Root --> A1["a"]
    Root --> B1["b"]

    A1 --> A2["b"]
    A1 --> C1["t"]

    B1 --> B2["a"]

    A2 --> A3["a"]
    A2 --> C2["t"]

    C1 --> C3["a"]
    C1 --> D1["l"]

    B2 --> B3["b"]

    B3 --> B4["a"]

    A3 --> A4["t"]

    C2 --> C4["n"]

    C3 --> C5["n"]
    D1 --> D2["a"]

    D2 --> D3["l"]