【Leetcode】642. Design Search Autocomplete System

题目地址：

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

要求设计一个搜索系统，其一开始会存储若干字符串以及它们对应的搜索次数，然后用户会连续输入字符，每次输入一个字符的时候，应当返回到该字符为止的字符串为前缀的、出现次数前$3$多的字符串，如果这样的字符串数量大于$3$了，则取字典序较小的。当用户输入'#'的时候，代表他输入完成了，此时要将他输入的整个字符串存入这个搜索系统，并且记录搜索次数加$1$。

思路是Trie + 最小堆 + 哈希表。我们可以将所有字符串存入Trie中，并且在每个Trie节点都预处理一下到当前节点为止的子串为前缀的所有字符串中，排名前$3$的是谁，这里可以用最小堆 + 哈希表来做，哈希表存每个字符串出现的次数，最小堆维护出现次数最高字典序最小的三个字符串。同时，我们需要用一个StringBuilder来存到当前为止，用户输入的字符串是什么。每次输入完一个字符，就对那个StringBuilder找答案；当输入了'#'的时候，就将StringBuilder对应的字符串$s$存入Trie并且次数加$1$。注意，这里每个Trie节点的最小堆也需要随着哈希表的计数做调整，所以最小堆要先删$s$再按照找前$3$的逻辑加回来。具体请看代码，代码如下：

import java.util.*;

public class AutocompleteSystem {
    
    class Trie {
        class Node {
            Node[] nexts;
            PriorityQueue<String> minHeap;
            
            public Node() {
                nexts = new Node[128];
                minHeap = new PriorityQueue<>(comp);
            }
        }
        
        Node root;
        
        public Trie(String[] sentences) {
            root = new Node();
            for (int i = 0; i < sentences.length; i++) {
                insert(sentences[i]);
            }
        }
        
        private void insert(String s) {
            Node cur = root;
            addToHeap(s, cur.minHeap);
            for (int i = 0; i < s.length(); i++) {
                char ch = s.charAt(i);
                if (cur.nexts[ch] == null) {
                    cur.nexts[ch] = new Node();
                }
                
                cur = cur.nexts[ch];
                // 在每个Trie节点都要尝试将s入堆
                addToHeap(s, cur.minHeap);
            }
        }
    }
    
    private Trie trie;
    // 记录每个字符串出现次数
    private Map<String, Integer> map;
    private Comparator<String> comp;
    // 记录当前用户输入了什么
    private StringBuilder sb;
    
    public AutocompleteSystem(String[] sentences, int[] times) {
        map = new HashMap<>();
        for (int i = 0; i < sentences.length; i++) {
            map.put(sentences[i], times[i]);
        }
        
        comp = (s1, s2) -> {
            int t1 = map.get(s1), t2 = map.get(s2);
            if (t1 != t2) {
                return Integer.compare(t1, t2);
            } else {
                return -s1.compareTo(s2);
            }
        };
        
        trie = new Trie(sentences);
        sb = new StringBuilder();
    }
    
    public List<String> input(char c) {
        List<String> res = new ArrayList<>();
        
        if (c == '#') {
            String s = sb.toString();
            // 先更新计数
            map.put(s, map.getOrDefault(s, 0) + 1);
            // 再把s插入到trie里
            trie.insert(s);
            // 刷新sb，以接受新输入
            sb = new StringBuilder();
            return res;
        }
        
        Trie.Node cur = trie.root;
        sb.append(c);
        
        for (int i = 0; i < sb.length(); i++) {
            char ch = sb.charAt(i);
            if (cur.nexts[ch] == null) {
                return res;
            }
            
            cur = cur.nexts[ch];
        }
        
        res.addAll(cur.minHeap);
        // 这里的比较器和堆的比较器是反序的
        res.sort(comp.reversed());
        return res;
    }
    
    private void addToHeap(String s, PriorityQueue<String> minHeap) {
    	// 为了防止s已经存在，所以要先将其删去，在哈希表计数更新了之后再加回来
        minHeap.remove(s);
        if (minHeap.size() < 3) {
            minHeap.offer(s);
            return;
        }
        
        String top = minHeap.peek();
        if (map.get(s) > map.get(top) || (map.get(s) == map.get(top) && s.compareTo(top) < 0)) {
            minHeap.poll();
            minHeap.offer(s);
        }
    }
}

初始化时间复杂度$O(nl)$，input时间复杂度$O(l)$，$n$是字符串数量，$l$是最长字符串长度，空间$O(nl)$。