425. Word Squares

Given a set of words (without duplicates), find all word squares you can build from them.

A sequence of words forms a valid word square if the kth row and column read the exact 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.

b a l l
a r e a
l e a d
l a d y

Note:

1. There are at least 1 and at most 1000 words.
2. All words will have the exact same length.
3. Word length is at least 1 and at most 5.
4. Each word contains only lowercase English alphabet a-z.

Example 1:

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

Output:
[
  [ "wall",
    "area",
    "lead",
    "lady"
  ],
  [ "ball",
    "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:
["abat","baba","atan","atal"]

Output:
[
  [ "baba",
    "abat",
    "baba",
    "atan"
  ],
  [ "baba",
    "abat",
    "baba",
    "atal"
  ]
]

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).

Solution:

We use a trie to store all the prefix in the word list.

In addition, in the TrieNode class, we add a List<String> member to store all the words start with this prefix.

In other words, at any node, we know all words from the word list that has the prefix to this node.

Therefore, we can use DFS to try all possible combinations. In each level, since we have a startWith list in the trie, we have already filtered out a large number of invalid candidates.



Code:

public class Solution {
    class TrieNode {
            public TrieNode[] link;
            public List<String> startWith;
            public TrieNode() {
                link = new TrieNode[26];
                startWith = new ArrayList<>();
            }
        }
        
    class Trie {
        public TrieNode root;
        
        public Trie(String[] words) {
            root = new TrieNode();
            
            for (String w : words) {
                TrieNode curr = root;
                for (char c : w.toCharArray()) {
                    if (curr.link[c - 'a'] == null) {
                        curr.link[c - 'a'] = new TrieNode();
                    }
                    curr = curr.link[c - 'a'];
                    curr.startWith.add(w);
                }
            }
        }
        
        public List<String> searchPrefix(String prefix) {
            TrieNode curr = root;
            for (char c : prefix.toCharArray()) {
                if (curr.link[c - 'a'] == null) {
                    return new ArrayList<String>();
                }
                curr = curr.link[c - 'a'];
            }
            return new ArrayList<String>(curr.startWith);
        }
    }
    
    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(words);
        List<String> list = new ArrayList<>();
        for (String w : words) {
            list.add(w);
            helper(trie, result, list);
            list.remove(list.size() - 1);
        }
        return result;
    }
    
    public void helper(Trie trie, List<List<String>> result, List<String> list) {
        if (list.size() == list.get(0).length()) {
            result.add(new ArrayList<>(list));
            return;
        }
        int index = list.size();
        StringBuilder prefix = new StringBuilder();
        for (String s : list) {
            prefix.append(s.charAt(index));
        }
        List<String> startWith = trie.searchPrefix(prefix.toString());
        for (String w : startWith) {
            list.add(w);
            helper(trie, result, list);
            list.remove(list.size() - 1);
        }
    }
}