Share my Java solution using Trie and DFS


  • 0
    M
    public class Solution {
        List<String> allWords = new ArrayList<>();
        public List<List<String>> wordSquares(String[] words) {
            Trie trie = new Trie('0');
            for (String word : words) {
                trie.addWord(word);
                allWords.add(word);
            }
            List<List<String>> list = new ArrayList<>();
            dfs(list, new ArrayList<>(), trie, words[0].length());
            return list;
        }
        public void dfs(List<List<String>> list, List<String> temp, Trie trie, int length) {
            if (temp.size()==length) {
                List<String> t = new ArrayList<>();
                for (String elem : temp)
                    t.add(elem);
                list.add(t);
                return;
            }
            int pos = temp.size();
            
            StringBuilder prefix = new StringBuilder();
            List<String> words = new ArrayList<String>();
            Trie current = trie;
            
            // fill in 'prefix' with the prefix of the next word
            // also update the Trie to the last char of prefix
            for (int i=0; i<temp.size(); i++) {
                prefix.append(temp.get(i).charAt(pos));
                
                // if we can't form a valid prefix, then return
                if (current.next[temp.get(i).charAt(pos)-'a'] == null)
                    return;
                current = current.next[temp.get(i).charAt(pos)-'a'];
            }
            
            // find all the possible words with given prefix
            if (pos > 0)
                findWords(words, current, prefix, length);
            else
                words = allWords;
            
            // dfs through all words
            for (String word : words) {
                temp.add(word);
                dfs(list, temp, trie, length);
                temp.remove(temp.size()-1);
            }
        }
        public void findWords(List<String> words, Trie trie, StringBuilder prefix, int length) {
            if (trie==null)
                return;
            if (prefix.length()==length) {
                words.add(prefix.toString());
                return;
            }
            
            // try all possible characters
            for (int i=0; i<26; i++) {
                if (trie.next[i]==null)
                    continue;
                prefix.append(trie.next[i].c);
                findWords(words, trie.next[i], prefix, length);
                prefix.deleteCharAt(prefix.length()-1);
            }
        }
        class Trie {
            char c;
            Trie[] next;
            public Trie(char c) {
                this.c = c;
                next = new Trie[26];
            }
            public void addWord(String word) {
                Trie current = this;
                for (int i=0; i<word.length(); i++) {
                    char ch = word.charAt(i);
                    if (current.next[ch-'a']==null)
                        current.next[ch-'a'] = new Trie(ch);
                    current = current.next[ch-'a'];
                }
            }
        }
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.