LeetCode 642: Design Search Autocomplete System

Note:

MANY DETAILS ARE MISSED:

1. If less than 3, return as many as it. So it could be NPE when queue.poll().getKey()

2. Use alphabetical order if times are same. So in the queu need to compare value first then compare string.

3. Logically it needs to set children first and move to child to add counts.

 

 

class AutocompleteSystem {
    class TrieNode {
        public Map<Character, TrieNode> children;
        public Map<String, Integer> counts;
        public TrieNode() {
            children = new HashMap<>();
            counts = new HashMap<>();
        }
    }
    
    private TrieNode _root;
    private StringBuilder _prefix;

    public AutocompleteSystem(String[] sentences, int[] times) {
        _root = new TrieNode();
        _prefix = new StringBuilder();
        for (int i = 0; i < sentences.length; i++) {
            addSentence(sentences[i], times[i]);
        }
    }
    
    private void addSentence(String sentence, int times) {
        TrieNode current = _root;
        for (char c : sentence.toCharArray()) {
            if (!current.children.containsKey(c)) {
                current.children.put(c, new TrieNode());
            }
            current = current.children.get(c);
            current.counts.put(sentence, current.counts.getOrDefault(sentence, 0) + times);
        }
    } 
    
    public List<String> input(char c) {
        if (c == '#') {
            addSentence(_prefix.toString(), 1);
            _prefix.setLength(0);
            return Collections.emptyList();
        }
        _prefix.append(c);
        TrieNode current = _root;
        for (char runner : _prefix.toString().toCharArray()) {
            if (!current.children.containsKey(runner)) {
                return Collections.emptyList();
            }
            current = current.children.get(runner);
        }
        
        PriorityQueue<Map.Entry<String, Integer>> queue = new PriorityQueue<>((a, b) -> a.getValue().compareTo(b.getValue()) == 0 ? a.getKey().compareTo(b.getKey()) : b.getValue() - a.getValue());
        queue.addAll(current.counts.entrySet());
        
        List<String> result = new ArrayList<>();
        for (int i = 0; i < 3; i++) {
            Map.Entry<String, Integer> entry = queue.poll();
            if (entry != null) {
                result.add(entry.getKey());
            }
        }
        return result;
    }
}

/**
 * Your AutocompleteSystem object will be instantiated and called as such:
 * AutocompleteSystem obj = new AutocompleteSystem(sentences, times);
 * List<String> param_1 = obj.input(c);
 */