LeetCode-1825. 求出 MK 平均值

地址：https://leetcode-cn.com/problems/finding-mk-average/
思路：https://leetcode-cn.com/problems/finding-mk-average/solution/c-san-ge-multiset-jian-dan-mo-ni-by-newh-y4q9/
维护 3 个 multiset：lower（保存最小的 k 个数）、middle（中间的数）、upper（保存最大的 k 个数）。
插入操作

• 如果 $num\le \max (lower)num\le max(lower)$，则在 lower 中插入 num
• 如果 $num\ge \min (upper)num\ge min(upper)$，则在 upper 中插入 num
• 否则，在 middle 中插入 num
  如果插入后，lower 或 upper 中的元素多于 k 个，则向 middle 中转移元素
  操作过程中维护 middle 的元素和 sum

删除操作

• 相当于滑动窗口，最左边的一个元素为要删除的元素，设删除的元素为 d
• d 一定存在于 lower 或 middle 或 upper 中的一个或多个集合中
• 选择一个删除即可
  如果删除后，lower 或 upper 中的元素少于 k 个，则向 middle 中索取元素
  操作过程中维护 middle 的元素和 sum

注意
对于multiset的删除函数erase()，不要直接用值删除lower.erase(d),会TLE，需要用迭代器的方式删除lower.erase(lower.rbegin())

class MKAverage {
    int m,k;
    queue<int> nums;
    multiset<int> lower, middle, upper;
    long long sum;
public:
    MKAverage(int m, int k) {
        this->m=m;
        this->k=k;
        sum=0;
    }
    
    void leftMove(multiset<int> &l, multiset<int> &r){
    	l.insert(*r.begin());
    	r.erase(r.begin());
	}
	
	void rightMove(multiset<int> &l, multiset<int> &r){
		r.insert(*l.rbegin());
		l.erase(--l.end());
	}
    
    void addElement(int num) {
        nums.push(num);
        if(!lower.empty() && *lower.rbegin() > num)	lower.insert(num);
		else if(!upper.empty() && *upper.begin() < num)	upper.insert(num);
		else middle.insert(num), sum+=num;
		while(lower.size() > k){
			sum+=*lower.rbegin(); rightMove(lower, middle);
		}
		while(upper.size() > k){
			sum+=*upper.begin(); leftMove(middle, upper);
		}
		
		if(nums.size() > m){	//滑动窗口的删除 
			int d = nums.front();	nums.pop();
			if(lower.find(d) != lower.end())	lower.erase(lower.find(d));
			else if(upper.find(d) != upper.end())	upper.erase(upper.find(d));
			else middle.erase(middle.find(d)), sum-=d;
		}
		
		if(nums.size() >= m){
			while(lower.size() < k){
				sum-=*middle.begin(); leftMove(lower, middle);
			}
			while(upper.size() < k){
				sum-=*middle.rbegin(); rightMove(middle, upper);
			}
		}
    }
    
    int calculateMKAverage() {
		if(nums.size() < m)	return -1;
		return sum/(m-2*k);
    }
};