【Leetcode】642. Design Search Autocomplete System

题目地址：

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

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

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

class AutocompleteSystem {
  using FSS = function<bool(string&, string&)>;
  struct Node {
    unordered_map<char, Node*> ne;
    // 这个堆是频率低的字符串优先，频率一样则字典序大的优先
    priority_queue<string, vector<string>, FSS> heap;
    unordered_set<string> inHeap;
    Node(FSS cmp) : heap(cmp) {}
  };

  struct Trie {
    Node* root;
    FSS cmp;

    Trie(FSS cmp) : cmp(cmp) { root = new Node(cmp); }

    void insert(string& s, int times, auto& mp) {
      mp[s] += times;
      Node* cur = root;
      addToHeap(s, cur);
      for (char ch : s) {
        if (!cur->ne[ch]) cur->ne[ch] = new Node(cmp);
        cur = cur->ne[ch];
        addToHeap(s, cur);
      }
    }

    void addToHeap(string& s, Node* node) {
      if (node->inHeap.find(s) == node->inHeap.end()) {
        node->heap.push(s);
        node->inHeap.insert(s);
        if (node->heap.size() > 3) {
          node->inHeap.erase(node->heap.top());
          node->heap.pop();
        }
      }
    }
  };

  Trie* trie;
  unordered_map<string, int> mp;
  FSS cmp;
  string input_str;

 public:
  AutocompleteSystem(vector<string>& ss, vector<int>& ts) {
    cmp = [&](string& s1, string& s2) {
      int t1 = mp[s1], t2 = mp[s2];
      return t1 != t2 ? t1 > t2 : s1 < s2;
    };

    trie = new Trie(cmp);
    for (int i = 0; i < ss.size(); ++i) trie->insert(ss[i], ts[i], mp);
  }

  vector<string> input(char c) {
    if (c == '#') {
      trie->insert(input_str, 1, mp);
      input_str.clear();
      return {};
    }

    input_str += c;
    Node* cur = trie->root;
    for (char ch : input_str) {
      if (!cur->ne[ch]) return {};
      cur = cur->ne[ch];
    }

    auto cp = cur->heap;
    vector<string> res;
    while (!cp.empty()) {
      res.push_back(cp.top());
      cp.pop();
    }
    sort(res.begin(), res.end(), cmp);
    return res;
  }
};

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

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