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

 1 public class Solution {
 2     public class TrieNode {
 3         TrieNode[] sons;
 4         List<String> startWith;
 5         public TrieNode() {
 6             this.sons = new TrieNode[26];
 7             this.startWith = new ArrayList();
 8         }
 9     }
10     
11     public class Trie {
12         TrieNode root;
13         public Trie() {
14             this.root = new TrieNode();
15         }
16         
17         public void insert(String word) {
18             char[] arr = word.toCharArray();
19             TrieNode node = root;
20             for (char c : arr) {
21                 if (node.sons[(int)(c-‘a‘)] == null) {
22                     node.sons[(int)(c-‘a‘)] = new TrieNode();
23                 }
24                 node.sons[(int)(c-‘a‘)].startWith.add(word);
25                 node = node.sons[(int)(c-‘a‘)];
26             }
27         }
28         
29         public List<String> findByPrefix(String prefix) {
30             char[] pre = prefix.toCharArray();
31             TrieNode node = root;
32             for (char c : pre) {
33                 if (node.sons[(int)(c-‘a‘)] == null) {
34                     return new ArrayList<String>();
35                 }
36                 node = node.sons[(int)(c-‘a‘)];
37             }
38             return node.startWith;
39         }
40     }
41     
42     public List<List<String>> wordSquares(String[] words) {
43         List<List<String>> res = new ArrayList<>();
44         if (words==null || words.length==0) return res;
45         int len = words[0].length();
46         
47         //build trie
48         Trie trie = new Trie();
49         for (String word : words) {
50             trie.insert(word);
51         }
52         
53         List<String> path = new ArrayList<String>();
54         for (String word : words) {
55             path.add(word);
56             helper(res, path, trie, len);
57             path.remove(path.size()-1);
58         }
59         return res;
60     }
61     
62     public void helper(List<List<String>> res, List<String> path, Trie trie, int len) {
63         if (path.size() == len) {
64             res.add(new ArrayList<String>(path));
65             return;
66         }
67         int pos = path.size();
68         StringBuilder prefix = new StringBuilder();
69         for (String str : path) {
70             prefix.append(str.charAt(pos));
71         }
72         List<String> nextPossible = trie.findByPrefix(prefix.toString());
73         for (String str : nextPossible) {
74             path.add(str);
75             helper(res, path, trie, len);
76             path.remove(path.size()-1);
77         }
78     }
79 }