Rainbow Subarray题解，SET动态维护中位数做法

SET 动态维护中位数

原理：原理很简单，就是开两个set（set1 和 set2）。对于一个集合，我们动态插入所有元素，使得两个set的大小差值不超过1（通常规定set2更大一点）且set2的最小值大于set1的最大值。这样我们就可以实时得到中位数的就是位于 $set2.begin()$ 的元素。
我们所要做的就是在这个过程中用set的借口来始终维护这两个性质即可。

代码有注释，变量名称也有标注

例题：

K. Rainbow Subarray

思路：由于和位置有关，所以把每个数字都剪去下标，然后用滑动窗口和set去动态维护中位数进行求解

#include <algorithm>
#include <cmath>
#include <cstring>
#include <iostream>
#include <set>
#include <vector>

#define int long long
const int N = 5e5;
int mapp[N + 2];
struct node {
    int v;
    node(int x) {
        v = x;
    }
    bool operator<(const node &t) const {
        return v < t.v;
    }
};
int n, m, sum1, sum2;  // 分别表示前半段 and 后半段的和
std::multiset<node> pre, pos;  // 分别表示前半段和后半段

int calc() {  // 这个函数与求解中位数无关，主要是为了解决把这一段所有的值调整到中位数上所需要的代价，（和具体题目有关）
    int median = pos.begin()->v;  // 这个就是中位数
    return (sum2 - median * (int)pos.size() + median * (int)pre.size() - sum1);
}

inline std::multiset<node>::iterator lst(std::multiset<node> &x) {
    std::multiset<node>::iterator res = x.end();  // 注意pre不能为空，这个一定要注意，否res--会出错
    res--;
    return res;
}
void adj() {  // 调整函数，目的是为了保证两段长度差不超过1且后半段较长
    if (pos.size() >= pre.size() + 2) {
        sum1 += pos.begin()->v;
        sum2 -= pos.begin()->v;
        pre.insert(pos.begin()->v);
        pos.erase(pos.begin());
        return;
    }
    if (pos.size() >= pre.size()) return;
    sum1 -= lst(pre)->v;
    sum2 += lst(pre)->v;
    pos.insert(lst(pre)->v);
    pre.erase(lst(pre));
}

void solve() {
    pre.clear(), pos.clear(), sum1 = 0, sum2 = 0;
    std::cin >> n >> m;
    for (int i = 1; i <= n; i++) {
        std::cin >> mapp[i];
        mapp[i] -= i;
    }

    //   要维护pos始终比pre长
    int ans = 0;
    for (int h = 1, t = 1; h <= n; h++) {
        node ls1(mapp[h]);
        if (mapp[h] >= pos.begin()->v) {
            pos.insert(mapp[h]);
            sum2 += mapp[h];
        } else {
            pre.insert(mapp[h]);
            sum1 += mapp[h];
        }
        adj();
        while (calc() > m && t <= h) {
            node ls2(mapp[t]);
            if ((pos.begin()->v) <= ls2.v) {
                sum2 -= mapp[t];
                pos.erase(pos.find(ls2));
                adj();
            } else {
                sum1 -= mapp[t];
                pre.erase(pre.find(ls2));
                adj();
            }
            t++;
        }
       
        ans = std::max(ans, h - t + 1);
    }
    std::cout << ans << std::endl;
}

signed main() {
    std::ios::sync_with_stdio(0);
    std::cin.tie(0);
    std::cout.tie(0);

    int t;
    std::cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}