425. Word Squares


Problem Description

The LeetCode problem asks for all possible word squares that can be formed from a given list of unique words. A word square is a special type of square grid of letters where every row and every column forms an identical word. This means that if you read the letters horizontally or vertically, the words should be the same. The sequence must satisfy the condition that the kth row's kth letter is the same as the kth column's kth letter, for every k from 0 up to the size of the largest row or column.

A crucial detail is that the same word can be used multiple times in building word squares, which greatly increases the number of potential combinations.

For example, with the words ["ball","area","lead","lady"], the following word square can be formed:

1ball
2area
3lead
4lady

Each row and column form the same words in this example. We are tasked with finding and returning all possible such squares that can be assembled given the input array of words.

Intuition

The intuition behind the solution involves several steps:

  1. Prefix Matching: A vital insight is that while building a word square, when we append a word, it determines the prefix that the next word must match. The progression is therefore by matching prefixes, word by word.

  2. Backtracking: An algorithmic technique that incrementally builds candidates to the solutions and abandons a candidate as soon as it determines this candidate cannot possibly lead to a valid word square.

  3. Trie Data Structure: This is a specialized tree-like data structure that's used for prefix matching. Each node stores links to other nodes - one for each possible alphabetic value. This structure efficiently retrieves all words that have a common prefix.

In the provided solution, we are using a Trie to store the words with a slight modification. Each Trie node keeps an array of indices, which is a list of all word indices from the original list that contain the prefix represented by the path from the root to that node.

Here's how we arrive at the solution:

  1. Building the Trie: We create a Trie and insert each word, keeping track of their indices within the original list.

  2. Searching the Trie: For building a word square, we start with the first word. For the next word to form a valid square, its prefix should match the second letter of the first word and so on. We use the Trie to look up all words with this prefix.

  3. Recursive Exploration with Backtracking: Next, we try to recursively add each of these words to the current word square in construction and check the next prefix that must be matched. If at any point, no words match the required prefix, we back out of the last step (we "backtrack") and try a different word. This process continues until:

    • We build a complete word square of the same size as the length of any word in the input list.
    • We exhaust all possibilities.

This solution method is efficient as it cuts down on the search space using the Trie structure for prefix look-ups and backtracking to avoid unnecessary exploration.

Solution Approach

The solution involves implementing a backtracking algorithm utilizing a Trie data structure to efficiently search for words that can form word squares. Let's walk through the implementation step by step:

  1. Trie Construction:

    • We define a Trie class with an initialization (__init__) method that creates 26 children (one for each letter of the alphabet) and an empty list v for keeping track of word indices.
    • The insert method adds a word to the Trie, character by character. For each character, we find the right child node to move to, or create a new Trie node if it does not exist. After reaching the end of a word, we append the word's index to the node's v list, capturing all indices where the path through the Trie matches the inserted word.
    • The search method follows the Trie nodes according to the characters of the given prefix, returning the list of word indices that contain the prefix if it exists, or an empty list otherwise.
  2. Backtracking with Trie:

    • In the wordSquares function of the Solution class, we first build a Trie from our words list, using the insert method for each word accompanied by its index.
    • We define an inner function dfs(t) for backtracking and building the word squares:
      • It takes the current partial word square t as an argument.
      • If t contains enough words to form a square (the length of t equals the length of a word), we add it to ans, our solution list.
      • Otherwise, for finding the next word, we look at the idxth characters of the words in t to form the next prefix we want to match.
      • We retrieve all indices of words with that prefix using the search method of the Trie.
      • We loop through these indices, adding each corresponding word to t and calling dfs(t) recursively before popping the word off t.
    • Finally, we start our search by iterating over each word in words and initiating a recursion with a word square starting with that word. The initial list [w] is passed to dfs to begin building the square.

The dfs function systematically constructs potential word squares, and the Trie enables fast lookups for words that can extend the partial squares by matching the required prefixes.

The wordSquares function ultimately returns ans, the collected list of all valid word squares, after exploring all possible combinations through recursive backtracking with the help of the Trie.

Example Walkthrough

Let's consider a smaller list of words to illustrate the solution approach: ["ab", "bc"]. The goal is to find all word squares that can be formed with these words. Here's how we would walk through the given solution approach:

  1. Trie Construction:

    • We create a Trie and insert each of the words with their indices. Our Trie will look like this after insertion:

      • Root
        • a (children: b - index [0])
        • b (children: c - index [1])
    • The Trie stores the index of "ab" under the branch 'a' -> 'b', and the index of "bc" under the branch 'b' -> 'c'.

  2. Backtracking with Trie:

    • We initialize ans as an empty list to collect valid word squares.
    • Starting with the word "ab", we consider this as the first word in our potential word square. This word determines that the second word in the square must start with 'b' because 'b' is the second character in "ab".
    • We then use the Trie to search for words beginning with 'b'. The Trie yields "bc", which matches the requirement.
    • We then form a partial word square ["ab", "bc"]. Since the number of words (2) is equal to the length of the words "ab" and "bc", we've successfully formed a valid word square and we add this to the ans list.

No other word squares can be formed using "ab" as the starting word. We would then repeat the process for "bc" as the starting word:

  1. Start with "bc" as the first word of the square.
  2. The second word must start with 'c'. However, our list does not have a word that starts with 'c', so no further word squares can be formed in this case.

The ans list now contains the single valid word square ["ab", "bc"]. We return this as the final result.

The example has been kept deliberately simple due to the limited number of words. In a scenario with a more extensive list of words, the Trie and backtracking process would be used more exhaustively to form all possible word squares.

Python Solution

1class Trie:
2    def __init__(self):
3        # Each Trie node contains an array of 26 Trie nodes, representing the 26 lowercase English letters
4        self.children = [None] * 26
5        # 'values' is a list that holds the indices of words corresponding to the node path
6        self.values = []
7
8    def insert(self, word, index):
9        # Insert a word into the Trie with corresponding index
10        node = self
11        for char in word:
12            char_index = ord(char) - ord('a')
13            if node.children[char_index] is None:
14                node.children[char_index] = Trie()
15            node = node.children[char_index]
16            # Append the index of the word at every character's node
17            node.values.append(index)
18
19    def search(self, prefix):
20        # Return a list of indices of words that start with the given prefix
21        node = self
22        for char in prefix:
23            char_index = ord(char) - ord('a')
24            if node.children[char_index] is None:
25                return []  # Prefix not found
26            node = node.children[char_index]
27        return node.values
28
29
30class Solution:
31    def wordSquares(self, words):
32        def dfs(square):
33            # Depth-first search to build word squares
34            if len(square) == len(words[0]):  # Base case: Square is complete
35                squares.append(square[:])  # Add a deep copy of the current square to results
36                return
37            # Get the current prefix to be matched from all words in the square
38            idx = len(square)
39            prefix = [word[idx] for word in square]
40            # Find all words in the Trie with the current prefix
41            indices = trie.search(''.join(prefix))
42            for index in indices:
43                square.append(words[index])  # Add the matching word to the current square
44                dfs(square)                 # Continue to build the square recursively
45                square.pop()                # Backtrack to try another word
46
47        trie = Trie()
48        squares = []
49        # Insert all words into the Trie along with their respective indices
50        for i, word in enumerate(words):
51            trie.insert(word, i)
52      
53        # Initialize the depth-first search with each word as a starting point
54        for word in words:
55            dfs([word])
56      
57        return squares
58

Java Solution

1import java.util.ArrayList;
2import java.util.Collections;
3import java.util.List;
4
5// Trie data structure to store words and their indices
6class Trie {
7    Trie[] children = new Trie[26];
8    List<Integer> indices = new ArrayList<>();
9
10    // Insert a word into the Trie, storing indices to trace the word in an array
11    void insert(String word, int index) {
12        Trie node = this;
13        for (char c : word.toCharArray()) {
14            int charIndex = c - 'a';
15            if (node.children[charIndex] == null) {
16                node.children[charIndex] = new Trie();
17            }
18            node = node.children[charIndex];
19            node.indices.add(index);
20        }
21    }
22
23    // Search for all words in the Trie that starts with a given prefix
24    List<Integer> search(String prefix) {
25        Trie node = this;
26        for (char c : prefix.toCharArray()) {
27            int charIndex = c - 'a';
28            if (node.children[charIndex] == null) {
29                return Collections.emptyList();
30            }
31            node = node.children[charIndex];
32        }
33        return node.indices;
34    }
35}
36
37public class Solution {
38    private Trie trie = new Trie();
39    private String[] wordsArray;
40    private List<List<String>> wordSquaresResult = new ArrayList<>();
41
42    public List<List<String>> wordSquares(String[] words) {
43        this.wordsArray = words;
44        for (int i = 0; i < words.length; i++) {
45            trie.insert(words[i], i);
46        }
47
48        List<String> wordSquare = new ArrayList<>();
49        for (String word : words) {
50            wordSquare.add(word);
51            depthFirstSearch(wordSquare);
52            wordSquare.remove(wordSquare.size() - 1);
53        }
54        return wordSquaresResult;
55    }
56
57    private void depthFirstSearch(List<String> wordSquare) {
58        if (wordSquare.size() == wordsArray[0].length()) {
59            wordSquaresResult.add(new ArrayList<>(wordSquare));
60            return;
61        }
62
63        int currentWordLength = wordSquare.size();
64        StringBuilder prefix = new StringBuilder();
65        for (String currentWord : wordSquare) {
66            prefix.append(currentWord.charAt(currentWordLength));
67        }
68
69        List<Integer> wordIndices = trie.search(prefix.toString());
70        for (int index : wordIndices) {
71            wordSquare.add(wordsArray[index]);
72            depthFirstSearch(wordSquare);
73            wordSquare.remove(wordSquare.size() - 1);
74        }
75    }
76}
77

C++ Solution

1#include <vector>
2#include <string>
3#include <unordered_map>
4using namespace std;
5
6// Definition for a Trie node
7class Trie {
8public:
9    unordered_map<char, Trie*> children;
10    vector<int> indices; // Stores indices of words to trace them in an array
11
12    // Method to insert a word into the Trie
13    void insert(const string& word, int index) {
14        Trie* node = this;
15        for (char c : word) {
16            if (node->children.find(c) == node->children.end()) {
17                node->children[c] = new Trie();
18            }
19            node = node->children[c];
20            node->indices.push_back(index);
21        }
22    }
23
24    // Method to search for all words in the Trie that start with a given prefix
25    const vector<int>& search(const string& prefix) const {
26        const Trie* node = this;
27        for (char c : prefix) {
28            if (node->children.find(c) == node->children.end()) {
29                return {}; // Return an empty vector if prefix not found
30            }
31            node = node->children.at(c);
32        }
33        return node->indices;
34    }
35};
36
37class Solution {
38private:
39    Trie trie; // Trie root to store words
40    vector<string> wordsArray; // Array to reference words by indices
41    vector<vector<string>> wordSquaresResult; // Stores the result of word squares
42
43    // Helper method to perform depth-first search to find word squares
44    void depthFirstSearch(vector<string>& wordSquare) {
45        if (wordSquare.size() == wordsArray[0].size()) { 
46            wordSquaresResult.push_back(wordSquare); // Found a valid word square
47            return;
48        }
49
50        // Build the prefix to search for in the Trie
51        string prefix;
52        int currentWordLength = wordSquare.size();
53        for (const string& currentWord : wordSquare) {
54            prefix += currentWord[currentWordLength];
55        }
56
57        // Get the indices of words with the current prefix
58        const vector<int>& indices = trie.search(prefix);
59        for (int index : indices) {
60            wordSquare.push_back(wordsArray[index]);
61            depthFirstSearch(wordSquare); // Recurse to build word square
62            wordSquare.pop_back(); // Backtrack
63        }
64    }
65
66public:
67    // Main method to find all word squares from a given list of words
68    vector<vector<string>> wordSquares(vector<string>& words) {
69        wordsArray = words; // Initialize the words array
70        for (int i = 0; i < words.size(); ++i) {
71            trie.insert(words[i], i); // Insert words into the Trie along with indices
72        }
73      
74        vector<string> wordSquare; // Current word square being built
75        for (const string& word : words) {
76            wordSquare.push_back(word);
77            depthFirstSearch(wordSquare); // Start DFS from each word
78            wordSquare.pop_back();
79        }
80        return wordSquaresResult; // Return all found word squares
81    }
82};
83

Typescript Solution

1// Global trie node structure definition
2interface TrieNode {
3    children: (TrieNode | null)[];
4    indices: number[];
5}
6
7// Function to create a new Trie node
8function createTrieNode(): TrieNode {
9    return {
10        children: new Array(26).fill(null),
11        indices: []
12    };
13}
14
15// Global variables
16let trie: TrieNode = createTrieNode();
17let wordsArray: string[];
18let wordSquaresResult: string[][] = [];
19
20// Function to insert a word into the Trie, with indices to track the word's position
21function insert(word: string, index: number) {
22    let node: TrieNode = trie;
23    for (let char of word) {
24        const charIndex = char.charCodeAt(0) - 'a'.charCodeAt(0);
25        if (!node.children[charIndex]) {
26            node.children[charIndex] = createTrieNode();
27        }
28        node = node.children[charIndex]!;
29        node.indices.push(index);
30    }
31}
32
33// Function to search for all words in the Trie that start with a given prefix
34function search(prefix: string): number[] {
35    let node: TrieNode | null = trie;
36    for (let char of prefix) {
37        const charIndex = char.charCodeAt(0) - 'a'.charCodeAt(0);
38        if (!node.children[charIndex]) {
39            return [];
40        }
41        node = node.children[charIndex];
42    }
43    return node.indices;
44}
45
46// Function to populate word squares given an array of words
47function wordSquares(words: string[]) {
48    wordsArray = words;
49    for (let i = 0; i < words.length; i++) {
50        insert(words[i], i);
51    }
52
53    let wordSquare: string[] = [];
54    for (let word of words) {
55        wordSquare.push(word);
56        depthFirstSearch(wordSquare);
57        wordSquare.pop();
58    }
59    return wordSquaresResult;
60}
61
62// Depth-first search helper function to build word squares
63function depthFirstSearch(wordSquare: string[]) {
64    if (wordSquare.length === wordsArray[0].length) {
65        wordSquaresResult.push([...wordSquare]);
66        return;
67    }
68
69    const currentWordLength = wordSquare.length;
70    let prefix = '';
71  
72    for (let currentWord of wordSquare) {
73        prefix += currentWord[currentWordLength];
74    }
75
76    let wordIndices = search(prefix);
77    for (let index of wordIndices) {
78        wordSquare.push(wordsArray[index]);
79        depthFirstSearch(wordSquare);
80        wordSquare.pop();
81    }
82}
83

Time and Space Complexity

Time Complexity

The time complexity of the wordSquares method primarily depends on the number of words n, the length of each word l, and the branching factor at each node of the Trie (which is at most 26, the number of lowercase English letters).

First, constructing the Trie takes O(n*l), because we traverse each word of length l and insert it into the Trie.

For each word, we start a depth-first search (DFS):

  • In the worst case, there can be n choices for the first word, n for the second, and so on, making it O(n^l) without optimization.
  • The Trie lookup reduces the branching factor significantly. For each node during DFS, instead of trying n possibilities, we look at the children that match the prefix formed by the current word square, which is an O(l) operation since at each recursive level we collect at most l indexes represented by the value list v.
  • The depth of the DFS is equal to the length of the words l, since we're forming squares, meaning we stop when the number of words equals the length of the words.

Assuming k is the average number of indexes stored in the Trie node's value list v after searching with prefixes, the approximate time complexity becomes O(n * k^(l-1) * l). This is because we perform the Trie search for every prefix formed in the recursion after the first word is chosen, and each search costs O(l) (since we're constructing the prefix from the partial word square and then iterating through the v list of the Trie node).

Space Complexity

The space complexity of the code is dictated by:

  • The Trie storage, which takes O(26*m*l), where m is the average length of each word when considering all nodes (with a maximum of 26 since each node has up to 26 children).
  • The recursion call stack, which will grow to a depth of l, along with the temporary array t in each recursive call and the pref constructed at each node, totaling O(l^2).
  • The solution array ans, which stores completed word squares. The total number of squares cannot exceed n^l, and each square has l words of length l, making it O(n^l * l^2).

Since the space taken by the solution ans will outweigh the others for large inputs, the dominating term in the space complexity is O(n^l * l^2), representing the output space.

Combining these, the overall space complexity is O(26 * m * l + l^2 + n^l * l^2). Simplifying, since 26 and m are constants and less significant than n^l, we get O(n^l * l^2).


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫