Leetcode: Word Squares && Summary: Another Important Implementation of Trie(Retrieve all the words with a given Prefix)

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:
There are at least 1 and at most 1000 words.
All words will have the exact same length.
Word length is at least 1 and at most 5.
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).

Backtracking + Trie:    referred to https://discuss.leetcode.com/topic/63516/explained-my-java-solution-using-trie-126ms-16-16

A better approach is to check the validity of the word square while we build it.
Example: ["area","lead","wall","lady","ball"]
We know that the sequence contains 4 words because the length of each word is 4.
Every word can be the first word of the sequence, let's take "wall" for example.
Which word could be the second word? Must be a word start with "a" (therefore "area"), because it has to match the second letter of word "wall".
Which word could be the third word? Must be a word start with "le" (therefore "lead"), because it has to match the third letter of word "wall" and the third letter of word "area".
What about the last word? Must be a word start with "lad" (therefore "lady"). For the same reason above.

The picture below shows how the prefix are matched while building the sequence.

In order for this to work, we need to fast retrieve all the words with a given prefix. There could be 2 ways doing this:

1. Using a hashtable, key is prefix, value is a list of words with that prefix.
2. Trie, we store a list of words with the prefix on each trie node.

The implemented below uses Trie.

public class Solution {
    public class TrieNode {
        TrieNode[] sons;
        List<String> startWith;
        public TrieNode() {
            this.sons = new TrieNode[26];
            this.startWith = new ArrayList();
        }
    }
    
    public class Trie {
        TrieNode root;
        public Trie() {
            this.root = new TrieNode();
        }
        
        public void insert(String word) {
            char[] arr = word.toCharArray();
            TrieNode node = root;
            for (char c : arr) {
                if (node.sons[(int)(c-'a')] == null) {
                    node.sons[(int)(c-'a')] = new TrieNode();
                }
                node.sons[(int)(c-'a')].startWith.add(word);
                node = node.sons[(int)(c-'a')];
            }
        }
        
        public List<String> findByPrefix(String prefix) {
            char[] pre = prefix.toCharArray();
            TrieNode node = root;
            for (char c : pre) {
                if (node.sons[(int)(c-'a')] == null) {
                    return new ArrayList<String>();
                }
                node = node.sons[(int)(c-'a')];
            }
            return node.startWith;
        }
    }
    
    public List<List<String>> wordSquares(String[] words) {
        List<List<String>> res = new ArrayList<>();
        if (words==null || words.length==0) return res;
        int len = words[0].length();
        
        //build trie
        Trie trie = new Trie();
        for (String word : words) {
            trie.insert(word);
        }
        
        List<String> path = new ArrayList<String>();
        for (String word : words) {
            path.add(word);
            helper(res, path, trie, len);
            path.remove(path.size()-1);
        }
        return res;
    }
    
    public void helper(List<List<String>> res, List<String> path, Trie trie, int len) {
        if (path.size() == len) {
            res.add(new ArrayList<String>(path));
            return;
        }
        int pos = path.size();
        StringBuilder prefix = new StringBuilder();
        for (String str : path) {
            prefix.append(str.charAt(pos));
        }
        List<String> nextPossible = trie.findByPrefix(prefix.toString());
        for (String str : nextPossible) {
            path.add(str);
            helper(res, path, trie, len);
            path.remove(path.size()-1);
        }
    }
}