LeetCode 1825. 求出 MK 平均值（set + queue）

1. 题目

给你两个整数 m 和 k ，以及数据流形式的若干整数。
你需要实现一个数据结构，计算这个数据流的 MK 平均值 。

MK 平均值 按照如下步骤计算：

• 如果数据流中的整数少于 m 个，MK 平均值 为 -1 ，否则将数据流中最后 m 个元素拷贝到一个独立的容器中。
• 从这个容器中删除最小的 k 个数和最大的 k 个数。
• 计算剩余元素的平均值，并 向下取整到最近的整数 。

请你实现 MKAverage 类：

• MKAverage(int m, int k) 用一个空的数据流和两个整数 m 和 k 初始化 MKAverage 对象。
• void addElement(int num) 往数据流中插入一个新的元素 num 。
• int calculateMKAverage() 对当前的数据流计算并返回 MK 平均数 ，结果需 向下取整到最近的整数 。

示例 1：
输入：
["MKAverage", "addElement", "addElement", "calculateMKAverage", "addElement", "calculateMKAverage", "addElement", "addElement", "addElement", "calculateMKAverage"]
[[3, 1], [3], [1], [], [10], [], [5], [5], [5], []]
输出：
[null, null, null, -1, null, 3, null, null, null, 5]
解释：
MKAverage obj = new MKAverage(3, 1); 
obj.addElement(3);        // 当前元素为 [3]
obj.addElement(1);        // 当前元素为 [3,1]
obj.calculateMKAverage(); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
obj.addElement(10);       // 当前元素为 [3,1,10]
obj.calculateMKAverage(); // 最后 3 个元素为 [3,1,10]
                          // 删除最小以及最大的 1 个元素后，容器为 [3]
                          // [3] 的平均值等于 3/1 = 3 ，故返回 3
obj.addElement(5);        // 当前元素为 [3,1,10,5]
obj.addElement(5);        // 当前元素为 [3,1,10,5,5]
obj.addElement(5);        // 当前元素为 [3,1,10,5,5,5]
obj.calculateMKAverage(); // 最后 3 个元素为 [5,5,5]
                          // 删除最小以及最大的 1 个元素后，容器为 [5]
                          // [5] 的平均值等于 5/1 = 5 ，故返回 5
 
提示：
3 <= m <= 10^5
1 <= k*2 < m
1 <= num <= 10^5
addElement 与 calculateMKAverage 总操作次数不超过 10^5 次。

来源：力扣（LeetCode） 链接：https://leetcode-cn.com/problems/finding-mk-average
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

2. 解题

• 使用 3个 multiset 记录3段数据，同时queue记录数据，更新删除，只保留 m 个
• 记录 中间段的 sum

class MKAverage {
    int m, k;
    long long sum = 0;//中间的数的和
    queue<int> q;
    multiset<int> big, mid;
    multiset<int, greater<int>> small;
public:
    MKAverage(int m, int k) {
        this->m = m;
        this->k = k;
    }
    
    void addElement(int num) {
        if(q.size() < m)
        {
            q.push(num);
            big.insert(num);//插入大的里面
            process();//处理3个set，保持有序，以及控制数量
        }
        else//满了
        {
            int tp = q.front();
            q.pop();
            q.push(num);
            big.insert(num);//插入大的里面
            if(mid.find(tp) != mid.end())
            {
                sum -= tp;
                mid.erase(mid.find(tp));
            }
            else if(big.find(tp) != big.end())
                big.erase(big.find(tp));
            else // (small.find(tp) != small.end())
                small.erase(small.find(tp));
            process();//处理3个set，保持有序，以及控制数量
        }
    }
    
    void process()
    {   // 处理 3个 set
        while(big.size() > k)
        {   // big 满了， 小的 放到 small 里
            small.insert(*big.begin());
            big.erase(big.begin());
        }
        while(small.size() > k)
        {   // small 满了，大的放到 mid 里
            mid.insert(*small.begin());
            sum += *small.begin(); // mid 和 增加
            small.erase(small.begin());
        }
        if(big.size()+mid.size()+small.size() < m)
            return;//数量不够 m 个， 结束逻辑
        if(*big.begin() < *mid.rbegin())
        { // 大的最小的 小于 中间的最大的， 交换数值
            int v1 = *mid.rbegin(), v2 = *big.begin();
            sum += v2-v1;
            mid.erase(--mid.end());
            big.erase(big.begin());
            mid.insert(v2);
            big.insert(v1);
        }
        if(big.size() < k)
        {   // 大的数量不够，中间的挪一个最大的到 大的
            auto it = --mid.end();
            sum -= *it;
            big.insert(*it);
            mid.erase(it);
        }
        if(*small.begin() > *mid.begin())
        {   // 小的最大的 大于 中间最小的，交换数值
            int v1 = *mid.begin(), v2 = *small.begin();
            sum += v2-v1;
            mid.erase(mid.begin());
            small.erase(small.begin());
            mid.insert(v2);
            small.insert(v1);
        }
        if(small.size() < k)
        {   // 小的数量不够，中间挪一个最小的到 小的
            auto it = mid.begin();
            sum -= *it;
            small.insert(*it);
            mid.erase(it);
        }
    }
    int calculateMKAverage() {
        if(q.size() < m)
            return -1;
        return sum/(m-2*k);
    }
};

560 ms 156.9 MB C++

---

我的CSDN博客地址 https://michael.blog.csdn.net/

长按或扫码关注我的公众号（Michael阿明），一起加油、一起学习进步！