Leetcode.1825 求出 MK 平均值

题目链接

Leetcode.1825 求出 MK 平均值

题目描述

给你两个整数 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\ast 2<m$
• $1<=num<=10^5$
• addElement与 calculateMKAverage总操作次数不超过 $10^5$ 次。

分析：

本题需要维护一个大小为 m的队列，里面的元素都是最新的 m个元素。此外我们还要设法快速求 队列中的元素除去 最大的k个元素 和 最小的k个元素 剩下元素和的平均值（向下取整）。
我们可以使用 两个树状数组 和 一个队列 维护这个数据结构。

树状数组可以看这篇博客：树状数组

• 一个队列 q，里面只存放最新的 m个元素
• 一个树状数组 s，存放插入元素的和
• 一个数组数组 cnt，存放插入元素的次数

代码：

using LL = long long;
const int N = 1e5+10;

class MKAverage {
public:
    //观察数据范围 ，记录元素和的 s 数组可能会爆int，所以使用 long long
    //这里为了方便 统一使用 long long
    LL cnt[N],s[N];
    
    int k,m;
    queue<int> q;
    MKAverage(int m, int k) {
        this->m = m;
        this->k = k;
        //将 cnt 和 s 数组都初始化为 0
        memset(cnt,0,sizeof cnt);
        memset(s,0,sizeof s);
    }

    //返回 x 的最低位的1
    //比如 x = 10111100
    //就返回 100 也即是4
    int lowbit(int x){
        return x & -x;
    }
    //添加一个元素
    void add(LL c[],int idx,int x){
        for(int i = idx;i <= N;i += lowbit(i)) c[i] += x;
    }
    //求 1 ~ idx 的前缀和
    LL preSum(LL c[],int idx){
        LL ans = 0;
        for(LL i = idx;i;i -= lowbit(i)) ans += c[i];
        return ans;
    }
    
    void addElement(int x) {
       //插入元素x ， s 记录插入元素的和，cnt 记录插入元素的次数 即每次加1
        q.push(x);
        add(s,x,x);
        add(cnt,x,1);
        
        //如果大于 m，就弹出头部的一个元素
        if(q.size() > this->m){
            int y = q.front();
            //队列头出队
            q.pop();
            //如下操作相当于 在cnt 和 s 中删除这个y
            add(s,y,-y);
            add(cnt,y,-1);
        }
    }

    LL getSum(int t){
        int l = 1,r = N;
        //用二分的方式快速找到 第一个前缀和（这里是元素的出现次数）大于等于 t 的元素
        //换句话说 就是找到 前 t 个元素之和（元素从小到大排）
        while(l < r){
            int mid = (l + r) >> 1;
            if(preSum(cnt,mid) >= t) r = mid;
            else l = mid + 1;
        }
        //此时的 r 就是第一个 次数的前缀和  大于等于 t 的元素
        //后面的一项 (preSum(cnt,r) - t) * 1LL * r，代表多余的元素
        //例如 假设依次插入了 [2,4,4,4,5]的元素，此时t = 2
        //所以我们要找第 2 个元素，也就是 r = 4
        //但是此时的 preSum(s,4) 实际上等于 18，实际上是四个元素的前缀和
        //因为 4 出现了三次 ， 所以我们要减去多余的次数之和
        //preSum(cnt,4) = 4
        //preSum(cnt,4) - t(t = 2) = 2
        //所以实际前 t 个元素之和应该为 preSum(s,4) - (preSum(cnt,4) - t)*4 = 6
        //这里 *1LL 的原因是将其转为 long long 防止溢出
        LL ans = preSum(s,r) - (preSum(cnt,r) - t) * 1LL * r;
        return ans;
    }
    
    int calculateMKAverage() {
        if(q.size() < m) return -1;
        else{
             //getSum(t) 即返回前 t 个数的和
            //假设此时 m = 6, k = 2 q = [2,3,3,5,5,9]
            //观察得 答案应该是中间两个数和的平均值
            //sum = getSum(4) - getSum(2) = 3 + 5 = 8;
            //所以答案 ans = sum / (6 - 2 * 2) = 4
            LL ans = (getSum(m - k) - getSum(k)) / (m - 2 * k);
            return ans;
        }
    }
};

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