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

Solution

public class Solution {

    public class Node {
        public Node[] children;
        public Node() {
            children = new Node[26];
        }
    }

    Node root = new Node();
    public List<List<String>> wordSquares(String[] words) {
        for(String word: words) {
            Node iter = root;
            for(char c: word.toCharArray()) {
                if(iter.children[c-'a'] == null) iter.children[c-'a'] = new Node();
                iter = iter.children[c-'a'];
            }
        }

        List<List<String>> ans = new ArrayList<List<String>>();
        List<String> path = new ArrayList<String>();
        for(String word: words) {
            path.add(word);
            dfs(ans, path);    
            path.remove(path.size()-1);
        }

        return ans;
    }

    public void dfs(List<List<String>> ans, List<String> path) {

        if(path.size() == path.get(0).length()) {
            ans.add(new ArrayList<String>(path));
            return;
        }

        // construct all possible words
        StringBuilder sb = new StringBuilder();
        Node iter = root;
        int seq = path.size();
        for(int i = 0; i < path.size(); i++) {
            char c = path.get(i).charAt(seq);

            if(iter.children[c-'a'] == null) return;
            sb.append(c);
            iter = iter.children[c-'a'];
        }

        // get every possible words, add it to path, and dfs further
        List<String> possiblePath = getPath(iter);
        for(String postfix: possiblePath) {
            path.add(sb.toString() + postfix);
            dfs(ans, path);
            path.remove(path.size()-1);
        }
    }

    public List<String> getPath(Node node) {
        List<String> ans = new ArrayList<String>();
        for(int i = 0; i < 26; i++) {
            if(node.children[i] != null) {
                for(String path: getPath(node.children[i])) {
                    ans.add((char)('a'+i) + path);
                }
            }
        }
        if(ans.size() == 0) ans.add("");
        return ans;
    }
}