[Google] LeetCode 1825 Finding MK Average 大根堆+小根堆

You are given two integers, m and k, and a stream of integers. You are tasked to implement a data structure that calculates the MKAverage for the stream.

The MKAverage can be calculated using these steps:

• If the number of the elements in the stream is less than m you should consider the MKAverage to be -1. Otherwise, copy the last m elements of the stream to a separate container.
• Remove the smallest k elements and the largest k elements from the container. Calculate the average value for the rest of the elements rounded down to the nearest integer.

Implement the MKAverage class:

• MKAverage(int m, int k) Initializes the MKAverage object with an empty stream and the two integers m and k.
• void addElement(int num) Inserts a new element num into the stream.
• int calculateMKAverage() Calculates and returns the MKAverage for the current stream rounded down to the nearest integer.

Solution

需要维护最近的 \(m\) 个元素，以及计算中间 \(m-2k\) 个元素的平均值。

既然要求前 \(k\) 大，前 \(k\) 小，那么应该想到利用类似于优先队列的结构，但是优先队列不方便删除元素。所以我们利用 \(set\).

维护一个大根堆 \(maxheap\), 一个小根堆 \(minheap\), 以及另一个小根堆 \(midheap\) 用来维护中间元素。其中 \(maxheap\) 维护前 \(k\) 小，\(minheap\) 维护前 \(k\) 大。

点击查看代码

class MKAverage {
private:
    set<pair<int,int>, greater<>> maxheap;
    set<pair<int,int>>minheap;
    set<pair<int,int>>midheap;
    
    queue<int> q;
    int cap;
    int tp;
    int cnt;
    long long sum;
    int len;
public:
    MKAverage(int m, int k) {
        cap=m;tp=k;
        sum=(long long)0;
        cnt=0;
        len = m-2*k;
    }
    
    void addElement(int num) {
        cnt++; q.push(num);
        if(cnt>cap){
            pair<int,int> rmv = {q.front(), cnt-cap};
            q.pop();
            
            if(maxheap.find(rmv)!=maxheap.end()){
                maxheap.erase(rmv);
                maxheap.insert(*midheap.begin());
                sum-=(*midheap.begin()).first;
                midheap.erase(midheap.begin());
            }
            
            else if(minheap.find(rmv)!=minheap.end()){
                minheap.erase(rmv);
                minheap.insert(*midheap.rbegin());
                sum-=(*midheap.rbegin()).first;
                midheap.erase(*midheap.rbegin());
            }
            else{
                // in midheap
                sum-=rmv.first;
                midheap.erase(midheap.find(rmv));
            }
        }
        // insert
        maxheap.insert({num, cnt});
        // #maxheap>K --> insert into minheap
        if(maxheap.size()>tp){
            auto f = *maxheap.begin();
            maxheap.erase(maxheap.begin());
            minheap.insert(f);
            
            // #minheap>K --> insert into midheap
            if(minheap.size()>tp){
                auto w = *minheap.begin();
                minheap.erase(minheap.begin());
                
                midheap.insert(w);
                sum+=w.first;
            }
        }
    }
    
    int calculateMKAverage() {
        if(cnt<cap)return -1;
        return int(sum/len);
    }
};

/**
 * Your MKAverage object will be instantiated and called as such:
 * MKAverage* obj = new MKAverage(m, k);
 * obj->addElement(num);
 * int param_2 = obj->calculateMKAverage();
 */