[H平衡树] lc1825. 求出 MK 平均值(multiset平衡树+STL代码细节+周赛236_4)

1. 题目来源

链接：lc1825. 求出 MK 平均值

2. 题目解析

动态维护一段数的最大值、最小值、增加、删除操作，那么就是平衡树了。一般手写平衡树的话考虑 splay 树，但是很难，竞赛的时候常用手写。在 C++ 中可以采用 set、multiset 简化实现，达到相同的效果。

平衡树的题目没接触过，看着题解写的勉勉强强。代码做过注释了。

思路：

• 维护三个 multiset 平衡树，分别对应维护三个区间，即 [0,k)、[k,k+m-2k)、[k+m-2k,m)，这三个区间，顺便维护一个 sum 是平衡树中值的总和，方便计算平均值。
• vector 维护数据流，小于 m 个就不用管，等于 m 个就排序后初始化三个平衡树，大于 m 个就无脑往中间平衡树插入，插入后，中间区间可能向左右出头到左右区间，因为是有序的，即 m.l < L.r 则说明无序了，需要交换这两个值，同理中间区间与右边区间也是如此。画图就很清楚了。
• 若加入后数据流长度大于 m，则需要删除一个数。在此查找该数的过程一定不要使用 count()，而需要使用 find()，貌似前者常数很大，按理说关联式容器应该都是 $O(logn)$ 才是…但使用 count() 会 TLE。也需要注意在这种关联式容器及平衡树中进行查找时，尽量使用容器本身的 find() 函数，而不要使用 <algorithm> 中的 find() 函数。

---

• 时间复杂度：$O(logn)$ 可能是…常数蛮大的
• 空间复杂度：$O(n)$

---

代码：
使用 count() 直接 TLE…

typedef long long LL;

class MKAverage {
public:
    struct Node{
        multiset<int> s;     // 底层红黑树
        LL sum = 0;

        // 增删 O(logn)
        void insert(int x){
            s.insert(x);
            sum += x;
        }
        void remove(int x){
            s.erase(s.find(x));
            sum -= x;
        }

    }L, M, R;   // 用三个multiset维护数据流的三个区间，自带有序
                // 分别是前k个元素,中间m-2k个元素，后面k个元素

    vector<int> q;      // 数据流
    int m, k;

    MKAverage(int _m, int _k) {
        m = _m, k = _k;
    }

    void addElement(int num) {
        q.push_back(num);

        // 如果输入个数小于m，不用管
        if (q.size() < m) return ;

        // 如果输入个数等于m,先对q排序，再初始化L M R
        if (q.size() == m) {
            auto t = q;
            sort(t.begin(), t.end());

            for (int i = 0; i < k; i ++) L.insert(t[i]);
            for (int i = k; i < m - k; i ++) M.insert(t[i]);
            for (int i = m - k; i < m; i ++) R.insert(t[i]);
        }
         // 如果输入个数大于m,维护L M R
        else {
            M.insert(num);      // 先将num丢入M中,再看是否使 L<= M <= R 失序

            if (*M.s.begin() < *L.s.rbegin()) {   // 检查左边
                int x = *M.s.begin(), y = *L.s.rbegin();
                L.insert(x); L.remove(y);
                M.insert(y); M.remove(x);
            }
            if (*M.s.rbegin() > *R.s.begin()) {   // 检查右边
                int x = *M.s.rbegin(), y = *R.s.begin();
                R.insert(x); R.remove(y);
                M.insert(y); M.remove(x);
            }

            // 加入一个元素后，维护的三个区间总长度超过 m，要删去最早添加的元素
            num = q[q.size() - 1 - m];
            // 查找最早添加的元素在哪个区间，删去它
            if (M.s.count(num)) M.remove(num);
            else if (L.s.count(num)) {
                L.remove(num);
                int x = *M.s.begin();   // 保持 L 区间长度不变
                M.remove(x);
                L.insert(x);
            } else {
                R.remove(num);
                int x = *M.s.rbegin();  // 保持 R 区间长度不变
                M.remove(x);
                R.insert(x);
            }
        }
    }

    int calculateMKAverage() {
        // 直接输出中间区间的平均值即可
        if (q.size() < m) return -1;
        return M.sum / M.s.size();
    }


};

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

使用 find() 就可以 ac 了：

typedef long long LL;

class MKAverage {
public:
    struct Node{
        multiset<int> s;     // 底层红黑树
        LL sum = 0;

        // 增删 O(logn)
        void insert(int x){
            s.insert(x);
            sum += x;
        }
        void remove(int x){
            s.erase(s.find(x));
            sum -= x;
        }

    }L, M, R;   // 用三个multiset维护数据流的三个区间，自带有序
                // 分别是前k个元素,中间m-2k个元素，后面k个元素

    vector<int> q;      // 数据流
    int m, k;

    MKAverage(int _m, int _k) {
        m = _m, k = _k;
    }

    void addElement(int num) {
        q.push_back(num);

        // 如果输入个数小于m，不用管
        if (q.size() < m) return ;

        // 如果输入个数等于m,先对q排序，再初始化L M R
        if (q.size() == m) {
            auto t = q;
            sort(t.begin(), t.end());

            for (int i = 0; i < k; i ++) L.insert(t[i]);
            for (int i = k; i < m - k; i ++) M.insert(t[i]);
            for (int i = m - k; i < m; i ++) R.insert(t[i]);
        }
         // 如果输入个数大于m,维护L M R
        else {
            M.insert(num);      // 先将num丢入M中,再看是否使 L<= M <= R 失序

            if (*M.s.begin() < *L.s.rbegin()) {   // 检查左边
                int x = *M.s.begin(), y = *L.s.rbegin();
                L.insert(x); L.remove(y);
                M.insert(y); M.remove(x);
            }
            if (*M.s.rbegin() > *R.s.begin()) {   // 检查右边
                int x = *M.s.rbegin(), y = *R.s.begin();
                R.insert(x); R.remove(y);
                M.insert(y); M.remove(x);
            }

            // 加入一个元素后，维护的三个区间总长度超过 m，要删去最早添加的元素
            num = q[q.size() - 1 - m];
            // 查找最早添加的元素在哪个区间，删去它
            if (M.s.find(num) != M.s.end()) M.remove(num);
            else if (L.s.find(num) != L.s.end()) {
                L.remove(num);
                int x = *M.s.begin();   // 保持 L 区间长度不变
                M.remove(x);
                L.insert(x);
            } else {
                R.remove(num);
                int x = *M.s.rbegin();  // 保持 R 区间长度不变
                M.remove(x);
                R.insert(x);
            }
        }
    }

    int calculateMKAverage() {
        // 直接输出中间区间的平均值即可
        if (q.size() < m) return -1;
        return M.sum / M.s.size();
    }
};

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