【Leetcode】1825. Finding MK Average

题目地址：

https://leetcode.com/problems/finding-mk-average/description/

要求设计一个数据结构，实现：
1、以两个正整数$m,k$初始化，$m>2k$；
2、添加一个元素；
3、计算最后添加的$m$个数里，去掉最大的$k$个和最小的$k$个的情况下的平均数。如果总共添加的数不足$m$个，则返回$-1$。

用一个map和一个队列维护最后的$m$个数。添加数的时候，只保留最后添加的$m$个数（即队列size大于$m$的话则pop队头，并从map中删掉这个数$1$次）。求平均数需要总和，最大的$k$个数的和，最小的$k$个数的和，总和可以直接开一个变量维护，最大的$k$个数可以从左到右遍历map求$k$个数的和，最小的$k$个数同理。代码如下：

class MKAverage {
 public:
  queue<int> q;
  map<int, int> mp;
  int m, k;
  long sum;

  MKAverage(int m, int k) {
    this->m = m;
    this->k = k;
    sum = 0;
  }

  void addElement(int x) {
    q.push(x);
    mp[x]++;
    sum += x;
	
	// 总数超过m了，删掉最先加入的数
    if (q.size() > m) {
      if (!--mp[q.front()]) mp.erase(q.front());
      sum -= q.front();
      q.pop();
    }
  }

  int calculateMKAverage() {
    if (q.size() < m) return -1;
    int n = k, minv = 0, maxv = 0;
    auto it = mp.begin();
    while (n > 0) {
      minv += it->first * min(it->second, n);
      n -= min(it->second, n);
      if (n) it++;
    }

    auto itr = mp.rbegin();
    n = k;
    while (n > 0) {
      maxv += itr->first * min(itr->second, n);
      n -= min(itr->second, n);
      if (n) itr++;
    }

    return (sum - minv - maxv) / (m - 2 * k);
  }
};

初始化时间复杂度$O(1)$，添加元素$O(\log m)$，求平均数$O(k)$，空间$O(m)$。