Explained. My Java solution using Trie [126ms 16/16]


  • 83
    L

    My first approach is brute-force, try every possible word sequences, and use the solution of Problem 422 (https://leetcode.com/problems/valid-word-square/) to check each sequence. This solution is straightforward, but too slow (TLE).

    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.

    0_1476809138708_wordsquare.png

    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 {
        class TrieNode {
            List<String> startWith;
            TrieNode[] children;
    
            TrieNode() {
                startWith = new ArrayList<>();
                children = new TrieNode[26];
            }
        }
    
        class Trie {
            TrieNode root;
    
            Trie(String[] words) {
                root = new TrieNode();
                for (String w : words) {
                    TrieNode cur = root;
                    for (char ch : w.toCharArray()) {
                        int idx = ch - 'a';
                        if (cur.children[idx] == null)
                            cur.children[idx] = new TrieNode();
                        cur.children[idx].startWith.add(w);
                        cur = cur.children[idx];
                    }
                }
            }
    
            List<String> findByPrefix(String prefix) {
                List<String> ans = new ArrayList<>();
                TrieNode cur = root;
                for (char ch : prefix.toCharArray()) {
                    int idx = ch - 'a';
                    if (cur.children[idx] == null)
                        return ans;
    
                    cur = cur.children[idx];
                }
                ans.addAll(cur.startWith);
                return ans;
            }
        }
    
        public List<List<String>> wordSquares(String[] words) {
            List<List<String>> ans = new ArrayList<>();
            if (words == null || words.length == 0)
                return ans;
            int len = words[0].length();
            Trie trie = new Trie(words);
            List<String> ansBuilder = new ArrayList<>();
            for (String w : words) {
                ansBuilder.add(w);
                search(len, trie, ans, ansBuilder);
                ansBuilder.remove(ansBuilder.size() - 1);
            }
    
            return ans;
        }
    
        private void search(int len, Trie tr, List<List<String>> ans,
                List<String> ansBuilder) {
            if (ansBuilder.size() == len) {
                ans.add(new ArrayList<>(ansBuilder));
                return;
            }
    
            int idx = ansBuilder.size();
            StringBuilder prefixBuilder = new StringBuilder();
            for (String s : ansBuilder)
                prefixBuilder.append(s.charAt(idx));
            List<String> startWith = tr.findByPrefix(prefixBuilder.toString());
            for (String sw : startWith) {
                ansBuilder.add(sw);
                search(len, tr, ans, ansBuilder);
                ansBuilder.remove(ansBuilder.size() - 1);
            }
        }
    }
    

  • 0
    Y

    The parameter "n" in search function can be left out.


  • 0
    L

    @yinan7 Thanks for pointing this out. I have updated the code accordingly.


  • 3
    S

    Is this method's space complexity huge? (startWith for each TrieNode)
    Storing the index of word in words may be better.


  • 0
    V
    This post is deleted!

  • 0
    F

    Concise and good solution with a list of start with string.


  • 0
    I

    What would be the time complexity of it? I am not sure about the search method especially.


  • 1
    C

    Hi, when you do this
    "List<String> startWith = tr.findByPrefix(prefixBuilder.toString());"

    is it possible that you could find the string you already added to the ans list?


  • 1
    S

    I think it will add duplicate strings to the ans. The test cases do not cover those corner cases yet.


  • 0
    C

    @shuoshankou Yeah. I tried something like ["aaaa", ......], and leetcode gives an answer with all "aaaa". So I'm a little confused. Is it allowed for this kind of situation?


  • 0

    @CBK Yes, the example below the question actually has duplicates.


  • 14

    Nice solution.

    One pic to help understand the Trie structure.

    The only difference between the trie here and the normal trie is that we hold one more list of all the words which have the prefix (from the root char to the current node char).

    alt text


  • 1
    L

    @LoseAGirlDoesNotMeanLoseASoul
    thanks for the explanation and the visualization.
    for folks who thought keeping a list of words would cost lots of memory: Java uses string pool, see details here: http://stackoverflow.com/questions/2486191/what-is-the-java-string-pool-and-how-is-s-different-from-new-strings


  • 0
    A

    Thanks for the nice solution. what is the time complexity?


  • 0
    A

    If a word starts with "aa", it may be added twice.


  • 1
    Y

    This turns out to be my Google onsite question...


  • 0
    R

    Just got asked this question by google on the phone. That was rough :)
    Actually, it was even harder than this one because my interviewer said that the words don't have to be the same length.


  • 0
    This post is deleted!

  • 1
    X

    As pointed out by @lzb700m, java has optimization over the string pool. I do not know whether there is an equivalent mechanism for C++. Anyway, we can use a standard declaration of Trie tree with DFS method to obtain all the words with a certain prefix.

        struct TrieNode {
            bool hasWord = false;
            unordered_map<char, TrieNode*> next;
            TrieNode(){}
        };
        
        struct Trie {
            TrieNode *root = new TrieNode();
            
            Trie(vector<string> &words){
                for(auto &s : words){
                    TrieNode *node = root;
                    for(auto &c : s){
                        if(node->next.count(c) == 0)
                            node->next.emplace(c, new TrieNode());
                        node = node->next[c];
                    }
                    node->hasWord = true;
                }   
            }
            
            vector<string> findPrefix(string prefix){
                vector<string> result;
                TrieNode *node = root;
                for(auto &c : prefix){
                    if(node->next.count(c) == 0)
                        return {};
                    node = node->next[c];
                }
                dfs(result, node, prefix);
                return result;
            }
            
            void dfs(vector<string> &result, TrieNode *node, string &prefix){
                if(node->hasWord)
                    result.push_back(prefix);
                    
                for(auto p : node->next){
                    prefix.push_back(p.first);
                    dfs(result, p.second, prefix);
                    prefix.pop_back();
                }
            }
        };
    

  • 0
    G

    @rmn Thanks for the info. I wonder if the interviewer further explained if you should consider prefixing or suffixing the string with empty spaces to form the word square? (in the case of LC 422 it's suffixing only)
    (1) suffix space only:
    a b c
    b d _
    c _ _
    (2) prefix space only:
    _ _ b
    _ c d
    b d _


Log in to reply
 

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