Course4all
Join us on telegram
Home/dsa/Trie/Word Squares

Word Squares

Master this topic with zero to advance depth.

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. A sequence of words forms a valid word square if the kth row and kth column read the exact same string for each 0 <= k < max(numRows, numColumns).

Hard

Examples

Input: words = ["area","lead","wall","lady","ball"]
Output: [["wall","area","lead","lady"],["ball","area","lead","lady"]]
Approach 1

Level I: Brute Force Backtracking

Intuition

Try all possible word combinations recursively. At each step, check if the current partial square is valid. Very slow as it doesn't use the prefix property effectively.

⏱ O(N^L) where N is words, L is length.💾 O(L)
java
class Solution {
    public List<List<String>> wordSquares(String[] words) {
        List<List<String>> res = new ArrayList<>();
        if (words == null || words.length == 0) return res;
        int n = words[0].length();
        for (String w : words) {
            List<String> list = new ArrayList<>(); list.add(w);
            backtrack(res, list, words, n);
        }
        return res;
    }
    private void backtrack(List<List<String>> res, List<String> list, String[] words, int n) {
        if (list.size() == n) { res.add(new ArrayList<>(list)); return; }
        int idx = list.size();
        StringBuilder sb = new StringBuilder();
        for (String s : list) sb.append(s.charAt(idx));
        String prefix = sb.toString();
        for (String word : words) {
            if (word.startsWith(prefix)) {
                list.add(word); backtrack(res, list, words, n); list.remove(list.size()-1);
            }
        }
    }
}
Approach 2

Level III: Trie + Backtracking

Intuition

At each step k, we need to find words that start with a specific prefix. This prefix is formed by the characters at position k of the words already chosen. A Trie is the perfect data structure to quickly fetch all words sharing a common prefix.

⏱ O(N * 26^L)💾 O(Total Chars)
java
import java.util.*;
class Solution {
    class Node { Node[] children = new Node[26]; List<String> wordList = new ArrayList<>(); }
    public List<List<String>> wordSquares(String[] words) {
        Node root = new Node();
        for (String w : words) {
            Node curr = root;
            for (char c : w.toCharArray()) {
                int i = c - 'a';
                if (curr.children[i] == null) curr.children[i] = new Node();
                curr = curr.children[i];
                curr.wordList.add(w);
            }
        }
        List<List<String>> res = new ArrayList<>();
        int n = words[0].length();
        for (String w : words) {
            List<String> square = new ArrayList<>();
            square.add(w); backtracking(1, n, square, root, res);
        }
        return res;
    }
    private void backtracking(int step, int n, List<String> square, Node root, List<List<String>> res) {
        if (step == n) { res.add(new ArrayList<>(square)); return; }
        StringBuilder prefix = new StringBuilder();
        for (String s : square) prefix.append(s.charAt(step));
        Node curr = root;
        for (char c : prefix.toString().toCharArray()) {
            if (curr.children[c-'a'] == null) return;
            curr = curr.children[c-'a'];
        }
        for (String candidate : curr.wordList) {
            square.add(candidate); backtracking(step + 1, n, square, root, res); square.remove(square.size()-1);
        }
    }
}
← Previous QuestionDesign Search Autocomplete System
Next Question →Maximum XOR of Two Numbers in an Array