[LeetCode] 425. Word Squares

Problem

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

Reference:
https://leetcode.com/problems...

Solution - Trie+DFS

class Solution {
    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();
        List<String> temp = new ArrayList<>();
        Trie trie = new Trie(words);
        for (String word: words) {
            temp.add(word);
            dfs(trie, len, temp, res);
            temp.remove(temp.size()-1);
        }
        return res;
    }
    
    private void dfs(Trie trie, int len, List<String> temp, List<List<String>> res) {
        int size = temp.size();
        if (size == len) {
            res.add(new ArrayList<>(temp));
            return;
        }
        StringBuilder sb = new StringBuilder();
        for (String str: temp) {
            sb.append(str.charAt(size));
        }
        String prefix = sb.toString();
        List<String> words = trie.getWords(prefix);
        for (String word: words) {
            temp.add(word);
            dfs(trie, len, temp, res);
            temp.remove(temp.size()-1);
        }
    }
}

class Trie {
    TrieNode root;
    public Trie(String[] strs) {
        root = new TrieNode();
        for (String str: strs) {
            TrieNode cur = root;
            for (char ch: str.toCharArray()) {
                int index = ch-'a';
                if (cur.children[index] == null) cur.children[index] = new TrieNode();
                cur.children[index].words.add(str);
                cur = cur.children[index];
            }
        }
    }
    
    public List<String> getWords(String prefix) {
        List<String> res = new ArrayList<>();
        TrieNode cur = root;
        for (char ch: prefix.toCharArray()) {
            int index = ch-'a';
            if (cur.children[index] == null) return res;
            cur = cur.children[index];
        }
        res.addAll(cur.words);
        return res;
    }
}

class TrieNode {
    TrieNode[] children;
    List<String> words;
    public TrieNode() {
        children = new TrieNode[26];
        words = new ArrayList<>();
    }
}