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

Very interesting problem. Let's explain by an example, consider we have 3 words in the list and we are trying to find the forth word. Based on the definition of word square, we know the prefix of the forth word should be the same as the forth character of each existing word in the list. e.g., :

wal | l
are  | a
lea  | d


Now all we need to do is find all words with prefix "lad", which is backtracking. A better way to search prefix is to use trie (no surprise). And that's it.

public class WordSquares {


    public List<List<String>> wordSquares(String[] words) {
        List<List<String>> rst = new ArrayList<>();
        if (words.length == 0) {
            return rst;
        }
        Trie trie = new Trie();
        for (String word : words) {
            trie.insert(word);
        }
        List<String> curr = new ArrayList<>();
        int len = words[0].length();
        for (String word : words) {
            curr.add(word);
            search(rst, trie, curr, len);
            curr.remove(curr.size() - 1);
        }
        return rst;
    }

    private void search(List<List<String>> rst, Trie trie, List<String> curr, int len) {
        if (curr.size() == len) {
            rst.add(new ArrayList<>(curr));
            return;
        }
        String prefix = "";
        int index = curr.size();
        for (String word : curr) {
            prefix += word.charAt(index);
        }
        List<String> startsWith = trie.prefixWith(prefix);
        for (String next : startsWith) {
            curr.add(next);
            search(rst, trie, curr, len);
            curr.remove(curr.size() - 1);
        }
    }

    private class Trie {
        TrieNode root;
        public Trie() {
            root = new TrieNode();
        }

        public void insert(String word) {
            TrieNode node = root;
            for (int i = 0; i < word.length(); i++) {
                char c = word.charAt(i);
                if (node.children[c - 'a'] == null) {
                    node.children[c - 'a'] = new TrieNode(c);
                }
                node = node.children[c - 'a'];
            }
            node.isWord = true;
        }

        public List<String> prefixWith(String prefix) {
            List<String> rst = new ArrayList<>();
            TrieNode node = root;
            for (int i = 0; i < prefix.length(); i++) {
                char c = prefix.charAt(i);
                if (node.children[c - 'a'] == null) {
                    return rst;
                }
                node = node.children[c - 'a'];
            }
            return node.prefixWith(prefix);
        }
    }





    private class TrieNode {
        char c;
        TrieNode[] children;
        boolean isWord;


        public TrieNode() {
            children = new TrieNode[26];
            isWord = false;
        }

        public TrieNode(char c) {
            this();
            this.c = c;
        }

        public List<String> prefixWith(String prefix) {
            List<String> rst = new ArrayList<>();
            if (isWord) {
                rst.add(prefix);
            }
            for (TrieNode child : children) {
                if (child != null) {
                    rst.addAll(child.prefixWith(prefix + child.c));
                }
            }
            return rst;
        }
    }
}