Saved from https://leetcode.com/problems/design-search-autocomplete-system/submissions/

lc642

class AutocompleteSystem {
    private static final int MAX_RESULT = 3;

    private final TrieNode root;
    private final StringBuilder currInput;
    private TrieNode prevNode;

    // Trie + Min heap
    // Time: O(N), N = number of sentences
    public AutocompleteSystem(String[] sentences, int[] times) {
      currInput = new StringBuilder();
      root = new TrieNode();
      prevNode = root;
      for (int i = 0; i < times.length; i++) {
        updateCountInTrie(root, sentences[i], times[i]);
      }
    }

    // Time: O(L * log3), L = the number of sentences in the sub-trie
    public List<String> input(char c) {
      if (c == '#') {
        updateCountInTrie(root, currInput.toString(), 1);
        currInput.setLength(0);
        prevNode = root;
        return Collections.emptyList();
      }
      // 1. record the input c
      // 2. if the sub-trie exist, return recommended sentences
      // 3. otherwise return
      currInput.append(c);
      if (prevNode == null || !prevNode.children.containsKey(c)) {
        prevNode = null;
        return Collections.emptyList();
      }
      TrieNode currNode = prevNode.children.get(c);
      PriorityQueue<TrieNode> queue = new PriorityQueue<>();
      if (currNode.sentence != null) {
        queue.add(currNode);
      }
      findTopWords(currNode, queue);
      prevNode = currNode;
      return generateResult(queue);
    }

    private List<String> generateResult(PriorityQueue<TrieNode> queue) {
      List<String> result = new ArrayList<>();
      while (!queue.isEmpty()) {
        result.add(queue.poll().sentence);
      }
      Collections.reverse(result);
      return result;
    }

    // Time: O(L*log(MAX_RESULT)), L = the number of sentences in the sub-trie
    private void findTopWords(TrieNode currNode, PriorityQueue<TrieNode> queue) {
      for (Map.Entry<Character, TrieNode> child : currNode.children.entrySet()) {
        // decide if add the sentence to the p-queue (compare child node's count with p-queue's min count)
        if (child.getValue().count > 0 && (queue.size() < MAX_RESULT || child.getValue().compareTo(queue.peek()) > 0)) {
          Objects.requireNonNull(child.getValue().sentence);
          queue.add(child.getValue());
          if (queue.size() > MAX_RESULT) {
            queue.poll();
          }
        }
        // traverse down the tree
        findTopWords(child.getValue(), queue);
      }
    }

    // Time: O(k), k = the length of the word
    private void updateCountInTrie(TrieNode root, String word, int count) {
      TrieNode curr = root;
      for (int i = 0; i < word.length(); i++) {
        curr.children.computeIfAbsent(word.charAt(i), k -> new TrieNode());
        curr = curr.children.get(word.charAt(i));
      }
      curr.sentence = word;
      curr.count += count;
    }

    static class TrieNode implements Comparable<TrieNode> {
      String sentence;
      int count;
      Map<Character, TrieNode> children = new HashMap<>();

      @Override
      public int compareTo(TrieNode o) {
        return this.count != o.count ? this.count - o.count : o.sentence.compareTo(this.sentence);
      }
    }
}

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