Rainbow Subarray

Rainbow Subarray

滑窗，中位数
2023 icpc济南站

题意

给定长度为 \(n\) 的序列 \(a_1,a_2,\dots,a_n\)，可以执行至多 \(k\) 次操作，每次操作使某个元素加一或减一。
一个连续子数组 \(a_l,a_{l+1},\dots,a_r\) 是彩虹子数组，当且仅当：

[
a_{i+1} - a_i = 1 \quad (l \le i < r)
]

长度为 \(1\) 的子数组始终是彩虹子数组。
求操作后最长彩虹子数组的长度。

---

核心转化

一个长度为 \(m\) 的窗口 \([l,r]\) 要变成彩虹子数组，目标序列一定可以写成：

[
c+l,\ c+l+1,\ \dots,\ c+r
]

其中 \(c\) 是任意整数。
位置 \(i\) 的目标值是 \(c+i\)，所以该位置的修改代价为：

[
|a_i - (c+i)|
]

令：

[
b_i = a_i - i
]

则：

[
|a_i - (c+i)| = |(a_i - i) - c| = |b_i - c|
]

于是窗口 \([l,r]\) 的最小修改代价为：

[
\min_c \sum_{i=l}^{r} |b_i - c|
]

这是一个经典的绝对偏差和问题，最小值在 \(c\) 取窗口内 \(b_i\) 的中位数时取得。

因此：

窗口 \([l,r]\) 可以变成彩虹子数组，当且仅当窗口内所有 \(b_i\) 到其中位数的绝对距离和 \(\le k\)。

---

双指针可行性

固定左端点，向右扩展窗口时，窗口的最小修改代价不会减少。
因为多加入一个元素，绝对偏差和只会增加或不变。

所以可以使用双指针：

• 右指针 \(r\) 不断向右扩展；
• 当窗口代价超过 \(k\) 时，左指针 \(l\) 向右移动并删除元素；
• 过程中记录最大窗口长度。

每个元素最多被加入一次、删除一次，总复杂度 \(O(n \log n)\)。

---

数据结构：两个 multiset 对顶维护中位数

用两个 multiset 维护当前窗口内的元素：

• small：存放较小的一半（包含中位数）
• large：存放较大的一半

保持：

[
|small| = |large| \quad \text{或} \quad |small| = |large| + 1
]

这样中位数就是 small 中的最大值：

*small.rbegin()

同时维护：

• sums：small 中所有元素的和
• suml：large 中所有元素的和

---

O(1) 计算窗口最小代价

设中位数为 \(mid\)，窗口长度为 \(m\)，则：

[
cost = \sum |b_i - mid|
]

因为 small 中元素都 \(\le mid\)，large 中元素都 \(> mid\)，所以：

[
cost = \sum_{x \in small}(mid - x) + \sum_{x \in large}(x - mid)
]

化简得：

[
cost = mid \cdot |small| - sums + suml - mid \cdot |large|
]

只要知道 \(mid\)、\(|small|\)、\(|large|\)、\(sums\)、\(suml\)，就可以 \(O(1)\) 计算代价。

---

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

void solve() {
    int n; ll k; cin >> n >> k;
    vector<ll> a(n + 1, 0), b(n + 1, 0);
    for(int i = 1; i <= n; i ++) cin >> a[i];
    for(int i = 1; i <= n; i ++) b[i] = a[i] - i;

    multiset<ll> large, small;
    ll suml = 0, sums = 0;
    auto balance = [&]() {
        while(small.size() > large.size() + 1) {
            auto it = prev(small.end());
            ll x = *it;
            large.insert(x); small.erase(it);
            sums -= x, suml += x;
        }
        while(large.size() > small.size()) {
            auto it = large.begin();
            ll x = *it;
            small.insert(x); large.erase(it);
            suml -= x, sums += x;
        }
    };
    auto add = [&](ll x) {
        if(small.empty() || x <= *small.rbegin()) {
            small.insert(x);
            sums += x;
        }
        else {
            large.insert(x);
            suml += x;
        }
        balance();
    };
    auto delet = [&](ll x) {
        auto it = small.find(x);
        if(it != small.end()) {
            small.erase(it);
            sums -= x;
        }
        else {
            it = large.find(x);
            large.erase(it);
            suml -= x;
        }
        balance();
    };
    auto getmin =[&]() -> ll {
        ll mid = *prev(small.end());
        ll L = small.size(), R = large.size();
        ll res = mid * L - sums + suml - R * mid;
        return res;
    };

    int l = 1;
    int ans = 0;
    for(int r = 1; r <= n; r ++) {
        add(b[r]);
        while(l <= r && getmin() > k) {
            delet(b[l]);
            l ++;
        }
        ans = max(ans, r - l + 1);
    }
    cout << ans << "\n";
}

int main() {  
    ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    int _ = 1; cin >> _;
    while(_--) solve();
    return 0;
}