【10.4】队列-求出 MK 平均值_队列取平均值-CSDN博客

输入：
["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 总操作次数不超过 105 次。

二、解题思路

我们使用三个有序集合 s1、s2 和 s3 来分别存储最小的 k 个元素、中间的 m−2k个元素以及最大的 k个元素。同时，我们用 sum2 来记录 s2中所有元素的和，并用队列 q 来保存最近的 m个元素。

addElement 函数的实现：

插入新元素：
- 将新元素 num加入队列 q。
- 根据队列 q 的当前大小进行以下操作：
情况 1：队列 q 的元素数量小于或等于 m：
- 将新元素 num插入有序集合 s2 中，并更新 sum2=sum2+num。
- 如果队列 q 的元素数量恰好等于 m，则需要将 s2 中最小的 k 个元素移动到 s1，并将最大的 k 个元素移动到 s3，同时更新 sum2。
情况 2：队列 q 的元素数量等于 m+1：
- 如果 num小于 s1 的最大元素，则将 num 插入 s1，并将 s1 的最大元素移动到 s2，同时更新 sum2。
- 如果 num 大于 s3 的最小元素，则将 num 插入 s3，并将 s3 的最小元素移动到 s2，同时更新 sum2。
- 否则，将 num 直接插入 s2，并更新 sum2=sum2+num。
- 由于 s2 的元素数量会多出 1 个，因此需要从 s1 或 s3中移动元素以保持平衡。
删除旧元素：
- 从队列 q 的队头取出最早加入的元素 x。
- 如果 x在 s2 中，则直接从 s2 中删除，并更新 sum2。
- 如果 x在 s1中，则从 s1中删除，并将 s2 中最小的元素移动到s1。
- 如果 x 在 s3 中，则从 s3 中删除，并将 s2 中最大的元素移动到 s3。

calculateMKAverage 函数的实现：

如果队列 q 中的元素数量小于 m，则直接返回 −1。
否则，返回 ⌊sum2/(m−2k)⌋，即 s2 中元素的平均值。

三、代码实现

#include <iostream>
#include <queue>
#include <set>
using namespace std;

class MKAverage {
private:
    int m, k;                // m 是窗口大小，k 是需要去掉的最小和最大元素的数量
    queue<int> q;            // 队列，用于存储当前窗口中的元素
    multiset<int> s1, s2, s3; // s1 存储最小的 k 个元素，s2 存储中间的 m-2k 个元素，s3 存储最大的 k 个元素
    long long sum2;          // 记录 s2 中所有元素的和

public:
    MKAverage(int m, int k) : m(m), k(k) {
        sum2 = 0; // 初始化 sum2 为 0
    }

    void addElement(int num) {
        q.push(num); // 将新元素加入队列

        if (q.size() <= m) {
            // 如果队列中的元素数量小于等于 m，直接将新元素加入 s2
            s2.insert(num);
            sum2 += num;

            if (q.size() == m) {
                // 当队列中的元素数量等于 m 时，初始化 s1 和 s3
                while (s1.size() < k) {
                    s1.insert(*s2.begin()); // 将 s2 中最小的元素移动到 s1
                    sum2 -= *s2.begin();
                    s2.erase(s2.begin());
                }
                while (s3.size() < k) {
                    s3.insert(*s2.rbegin()); // 将 s2 中最大的元素移动到 s3
                    sum2 -= *s2.rbegin();
                    s2.erase(prev(s2.end()));
                }
            }
            return;
        }

        // 如果队列中的元素数量超过 m，需要调整 s1、s2 和 s3
        if (num < *s1.rbegin()) {
            // 如果新元素小于 s1 的最大元素，将其加入 s1，并将 s1 的最大元素移动到 s2
            s1.insert(num);
            s2.insert(*s1.rbegin());
            sum2 += *s1.rbegin();
            s1.erase(prev(s1.end()));
        } else if (num > *s3.begin()) {
            // 如果新元素大于 s3 的最小元素，将其加入 s3，并将 s3 的最小元素移动到 s2
            s3.insert(num);
            s2.insert(*s3.begin());
            sum2 += *s3.begin();
            s3.erase(s3.begin());
        } else {
            // 否则，直接将新元素加入 s2
            s2.insert(num);
            sum2 += num;
        }

        // 移除队列中最旧的元素
        int x = q.front();
        q.pop();

        if (s1.count(x) > 0) {
            // 如果 x 在 s1 中，将其从 s1 中移除，并将 s2 的最小元素移动到 s1
            s1.erase(s1.find(x));
            s1.insert(*s2.begin());
            sum2 -= *s2.begin();
            s2.erase(s2.begin());
        } else if (s3.count(x) > 0) {
            // 如果 x 在 s3 中，将其从 s3 中移除，并将 s2 的最大元素移动到 s3
            s3.erase(s3.find(x));
            s3.insert(*s2.rbegin());
            sum2 -= *s2.rbegin();
            s2.erase(prev(s2.end()));
        } else {
            // 如果 x 在 s2 中，直接将其从 s2 中移除
            s2.erase(s2.find(x));
            sum2 -= x;
        }
    }

    int calculateMKAverage() {
        if (q.size() < m) {
            // 如果队列中的元素数量小于 m，返回 -1
            return -1;
        }
        // 返回 s2 中元素的平均值
        return sum2 / (m - 2 * k);
    }
};

int main() {
    // 测试代码
    MKAverage mkAvg(3, 1); // 初始化窗口大小为 3，去掉 1 个最小和 1 个最大元素

    mkAvg.addElement(1);
    mkAvg.addElement(3);
    mkAvg.addElement(-1);

    cout << "MKAverage: " << mkAvg.calculateMKAverage() << endl; // 输出结果

    mkAvg.addElement(-3);
    mkAvg.addElement(5);
    mkAvg.addElement(3);

    cout << "MKAverage: " << mkAvg.calculateMKAverage() << endl; // 输出结果

    return 0;
}