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:

build a trie so that it's easier to find strings with prefix. calculate the next prefix with current list of strings, for example
"wall" => "a"
"wall", "area" => "le"
"wall", "area", "lead" => "lad"

class Solution {
    class TrieNode {
        TrieNode[] children = new TrieNode[26];
        boolean isWord = false;
    }
    
    public void insert(TrieNode root, String w) {
        for (int i = 0; i < w.length(); i ++) {
            char c = w.charAt(i);
            if (root.children[c - 'a'] == null) {
                root.children[c - 'a'] = new TrieNode();
            }
            root = root.children[c - 'a'];
        }
        root.isWord = true;
    }
    
    public List<String> search(TrieNode root, String prefix) {
        for (int i = 0; i < prefix.length(); i ++) {
            char c = prefix.charAt(i);
            if (root.children[c - 'a'] != null) {
                root = root.children[c - 'a'];
            } else {
                return new ArrayList();
            }
        }
        List<String> result = new ArrayList();
        traverse(root, new StringBuilder(prefix), result);
        return result;
    }
    
    private void traverse(TrieNode root, StringBuilder sb, List<String> result) {
        if (root.isWord) {
            result.add(sb.toString());
            return;
        }
        for (char c = 'a'; c <= 'z'; c ++) {
            if (root.children[c - 'a'] != null) {
                sb.append(c);
                traverse(root.children[c - 'a'], sb, result);
                sb.deleteCharAt(sb.length() - 1);
            }
        }
    }
    
    public List<List<String>> wordSquares(String[] words) {
        TrieNode root = new TrieNode();
        for (String w : words) {
            insert(root, w);
        }
        List<List<String>> result = new ArrayList();
        dfs(root, words[0].length(), new ArrayList(), result);
        return result;
    }
    
    private void dfs(TrieNode root, int len, List<String> curr, List<List<String>> result) {
        if (curr.size() == len) {
            result.add(new ArrayList(curr));
            return;
        }
        String prefix = calcPrefix(curr);
        List<String> possibles = search(root, prefix);
        for (String next : possibles) {
            curr.add(next);
            dfs(root, len, curr, result);
            curr.remove(curr.size() - 1);
        }
    }
    
    private String calcPrefix(List<String> curr) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < curr.size(); i ++) {
            String w = curr.get(i);
            sb.append(w.charAt(curr.size()));
        }
        return sb.toString();
    }
}