Java solution, Trie and PriorityQueue


  • 18

    Only thing more than a normal Trie is added a map of sentence to count in each of the Trie node to facilitate process of getting top 3 results.

    public class AutocompleteSystem {
        class TrieNode {
            Map<Character, TrieNode> children;
            Map<String, Integer> counts;
            boolean isWord;
            public TrieNode() {
                children = new HashMap<Character, TrieNode>();
                counts = new HashMap<String, Integer>();
                isWord = false;
            }
        }
        
        class Pair {
            String s;
            int c;
            public Pair(String s, int c) {
                this.s = s; this.c = c;
            }
        }
        
        TrieNode root;
        String prefix;
        
        
        public AutocompleteSystem(String[] sentences, int[] times) {
            root = new TrieNode();
            prefix = "";
            
            for (int i = 0; i < sentences.length; i++) {
                add(sentences[i], times[i]);
            }
        }
        
        private void add(String s, int count) {
            TrieNode curr = root;
            for (char c : s.toCharArray()) {
                TrieNode next = curr.children.get(c);
                if (next == null) {
                    next = new TrieNode();
                    curr.children.put(c, next);
                }
                curr = next;
                curr.counts.put(s, curr.counts.getOrDefault(s, 0) + count);
            }
            curr.isWord = true;
        }
        
        public List<String> input(char c) {
            if (c == '#') {
                add(prefix, 1);
                prefix = "";
                return new ArrayList<String>();
            }
            
            prefix = prefix + c;
            TrieNode curr = root;
            for (char cc : prefix.toCharArray()) {
                TrieNode next = curr.children.get(cc);
                if (next == null) {
                    return new ArrayList<String>();
                }
                curr = next;
            }
            
            PriorityQueue<Pair> pq = new PriorityQueue<>((a, b) -> (a.c == b.c ? a.s.compareTo(b.s) : b.c - a.c));
            for (String s : curr.counts.keySet()) {
                pq.add(new Pair(s, curr.counts.get(s)));
            }
    
            List<String> res = new ArrayList<String>();
            for (int i = 0; i < 3 && !pq.isEmpty(); i++) {
                res.add(pq.poll().s);
            }
            return res;
        }
    }
    

  • 0
    A

    Nice solution! Pasting my shortened and commented Python translation -

    class AutocompleteSystem(object):  
        def __init__(self, sentences, times):
            self.root = {}
            self.prefix = ""
            
            for sentence, freq in zip(sentences, times):
                self.insert(sentence, freq)
        
        def insert(self, sentence, freq): 
            node = self.root
            for letter in sentence:
                if letter not in node: 
                    node[letter] = {}
                    node[letter]['topx'] = collections.Counter()
    
                node = node[letter] #advance the pointer first, then increment the frequency counter for that letter pos. 
                node['topx'][sentence] += freq
                
        def input(self, c):
            if c == '#': 
                self.insert(self.prefix, 1) #whether its an existing sentence or a new sentence, we increment freq. by 1. 
                self.prefix = "" #reset
                return [] #not asked to return any suggestion if the user has terminated the input. 
    
            self.prefix = self.prefix + c
            node = self.root
            for cc in self.prefix:
                if cc not in node: 
                    return []
                node = node[cc]
    
            suggestions = zip(node['topx'].values(), node['topx'].keys())
            suggestions.sort(key = lambda x: (-x[0], x[1])) #sort by freq. desc, then by lex. asc 
            return map(lambda x: x[1], suggestions[:3]) 
    

  • 7

    Another Approach.

    public class AutocompleteSystem {
    
        private Map<String, Integer> map = new HashMap<>();
        private StringBuilder build = new StringBuilder ();
        private List<Map.Entry<String, Integer>> answers = new ArrayList<> ();
        
        public AutocompleteSystem(String[] sentences, int[] times) {
            for (int idx = 0; idx < sentences.length; idx ++) map.put (sentences [idx], times [idx]);
        }
        
        public List<String> input(char c) {
            List<String> ans = new ArrayList<>();
            if (c == '#') {  
                map.put (build.toString (), map.getOrDefault (build.toString (), 0) + 1); 
                build.setLength (0); 
                answers.clear(); 
            } else {
                build.append (c);
                if (build.length () == 1) {
                    for (Map.Entry<String, Integer> e : map.entrySet ())  
                        if (e.getKey ().startsWith (build.toString ())) answers.add (e);
                    
                    Collections.sort (answers,  (a, b) -> (a.getValue() == b.getValue()) ? a.getKey ().compareTo (b.getKey ()) : b.getValue() - a.getValue());
                } else {
                    for (Iterator <Map.Entry <String, Integer>> itr = answers.iterator (); itr.hasNext();) 
                        if (!itr.next().getKey().startsWith (build.toString ())) itr.remove ();
                }
                for (int idx = 0; idx < 3 && idx < answers.size(); idx ++) ans.add (answers.get (idx).getKey ());
            }
            return ans;
        }
    }
    
    

  • 0
    N

    @Himanshu like this solution, clean and simple


  • 1

    Maybe it's better to store the sentence only in its own node, instead of all its prefix nodes. Do a tree traverse to get children sentences won't increase the time complexity.

    public class AutocompleteSystem {
        private TrieNode root;
        private String input;
        private TrieNode currNode;
        
        public AutocompleteSystem(String[] sentences, int[] times) {
            root = new TrieNode("");
            input = "";
            currNode = root;
            for (int i = 0; i < sentences.length; i++) { //build the prefix tree
                root.add(sentences[i], 0, times[i]);
            }
        }
        
        public List<String> input(char c) {
            if(c == '#') {
                currNode.count++;
                currNode = root; //prepare for next input string
                input ="";
                return Collections.emptyList();
            }
            input = input+c;
            currNode.children.putIfAbsent(c, new TrieNode(input));
            currNode = currNode.children.get(c);
            return currNode.getTop3();
        }
        
        class TrieNode{
            HashMap<Character, TrieNode> children;
            String str;//sentence/prefix end here
            int count;//num of sentences end here;
            
            public TrieNode(String str) {
                children = new HashMap<>();
                this.str = str;
                count = 0;
            }
            
            public void add(String sentence, int i, int num) {
                if (i == sentence.length()) {
                    this.count += num;
                    return;
                }
                char ch = sentence.charAt(i);
                children.putIfAbsent(ch, new TrieNode(sentence.substring(0, i+1)));
                children.get(ch).add(sentence, i+1, num);
            }
            
            public List<String> getTop3(){
                PriorityQueue<SentenceData> pq = new PriorityQueue<>();
                preOrderTraverse(pq);
                LinkedList<String> res = new LinkedList<>();
                while(!pq.isEmpty()) res.addFirst(pq.poll().sentence);
                return res;
            }
            
            private void preOrderTraverse(PriorityQueue<SentenceData> pq) {
                if (count >0) {
                    pq.offer(new SentenceData(this.str, count)); //MinHeap
                    if (pq.size()>3) pq.poll();
                }                       
                for (TrieNode child : children.values()) {
                    child.preOrderTraverse(pq);
                }
            }
            
            class SentenceData implements Comparable<SentenceData>{
                String sentence;
                int count;
                public SentenceData(String s, int c){
                    this.sentence = s;
                    this.count = c;
                }     
                @Override
                public int compareTo(SentenceData other){
                    if (this.count == other.count) {
                        return other.sentence.compareTo(this.sentence);
                    }
                    return this.count-other.count; // Asending order
                }
            }
        }
    }
    

  • 0
    A

    @Himanshu Time complexity of the solution suffers though when you use hashmap and traverse each sentence every time, instead of using a trie.


  • 0

    @aditya74 I am not going through each sentence everytime. It's only done once after the '#' and later on the answers are filtered down. I understand this can be easily done with trie, dfsing the whole tree emanating from the start char and retaining the answer for future filteration.


  • 0
    O
    This post is deleted!

  • 0
        Trie root;
        public AutocompleteSystem(String[] sentences, int[] times) {
            root = new Trie();
            for(int i = 0; i < sentences.length; i++){
              insert(root, sentences[i], times[i]);  
            }
        }
        
        private void insert(Trie root, String sentence, int time){
           for(char c: sentence.toCharArray()){
             if(root.children[index(c)] == null){
               root.children[index(c)] = new Trie();  
             }  
             root = root.children[index(c)];  
           } 
           if(root.node == null){
               root.node = new Node(sentence, time); 
           } else {
               root.node.time += time;
           }
        }
        
        private int index(char c){
           return c == ' ' ? 26 : c - 'a';   
        }
        
        String cur = "";
        public List<String> input(char c) {
            List<String> res = new ArrayList();
            if(c == '#'){
                insert(root, cur, 1);
                cur = "";
            } else {
                cur += c;
                List<Node> list = lookup(root, cur);
                Collections.sort(list, (a, b) -> a.time == b.time ? a.sentence.compareTo(b.sentence) : b.time - a.time);
                for (int i = 0; i < Math.min(3, list.size()); i++)
                    res.add(list.get(i).sentence);
            }
            return res;
        }
        
        private List<Node> lookup(Trie root, String cur){
           List<Node> res = new ArrayList(); 
           for(char c: cur.toCharArray()){
              if(root.children[index(c)] == null) return res; 
              else root = root.children[index(c)]; 
           } 
           
           helper(root, res);
           return res;
        }
        
        private void helper(Trie root, List<Node> res){
           if(root == null) return; 
           if(root.node != null) res.add(root.node); 
           for(Trie t: root.children){
              helper(t, res); 
           } 
        }
    }
    
    class Node{
      public String sentence;
      public int time;
      public Node(String s, int t){
         sentence = s;
         time = t; 
      }  
    }
    
    class Trie{
      public Node node;
      public Trie[] children = new Trie[27]; 
    }

  • 1
    W

    @shawngao why don't you save a topK heap inside TrieNode instead of saving a hashmap?


  • 0

    @wminghao Because for heap pq, it is not easy to update its existing elements when calling add(string, count) after a search entry is fininshed.

    Also, assuming you could update heap, you actually push all the complexity to function add(string, count) to update _prefix.size() many heaps instead of processing a single heap each time a user types a single character by invoking input(char c). Since we already use Trie to optimize the prefix search, I guess the logic behind scene is to spread out the burden to each character typed and expecting that the possible matches will be quickly narrowed down as more characters are typed.

    But on a second thought ... Since autocomplete system should be a read heavy system (is it? ...), maybe it is preferable as you suggested to push all the update work in "add", which is the write query, and get all heaps ready for input, the read query.


  • 1

    @wminghao @zzg_zzm That makes sense. We can use a PriorityQueue inside the TrieNode instead of a HashMap. I thought of this during contest but as @zzg_zzm mentioned it is a little bit complicated to update contents in each queue when we want to add a new word, also every time we need to pop the first K elements out and put them back again when we want to generate the result.
    To save time, I chose HashMap approach. Remember to discuss with your interviewer if it is read heavy or write heavy and tell him/her pros/cons of different solutions, you will get a "STRONG HIRE" :)


  • 0
    W

    @zzg_zzm thanks. From the function description, the input string list don't change. So I assume it's a read only system. But as you said, if it can be changed, hopefully it won't be changed too frequently, like update daily, we can still use heap and build a new Trie every night.

    I don't see a strong reason of using a hashmap, inserting into heap everytime a user type something is really inefficient.


  • 0

    @shawngao why there is a isWord attribute of TrieNode? I didn't see why you use it. Could you explain?


  • 0
    K

    Could you express why you need to record count in node?


  • 0
    J

    We could add Map.Entry<String, Integer> to the PriorityQueue instead of using a Pair class, is there any other reason to do so?


  • 0
    W

    @shawngao Just feel that copies of complete sentences are stored at every TrieNode along its path down, which has pretty high memory overhead.


  • 0
    W

    @jidong_jack you may notice that Shawngao actually created new pairs class and add them to PriorityQueue (PQ). If you put Map.Entry from counts (hasmap) to PQ, I am not sure this is correct. If you create new Map.Entry and put PQ, it should be OK.


Log in to reply
 

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