P16318 [ICPC 2023 Jinan R] 彩虹子数组

思路

首先注意到一段公差为 $1$ 的等差数列等价于将所有值减掉其位置后的连续相等数。
所以先给每个 $a_i$ 减掉 $i$，这样就把原问题转化为求一段连续相等区间的代价。
此时对于一段区间成为彩虹子数组的代价即为所有数都向上或向下变成同一个数的代价。
则这个最终数显然是中位数。
求出中位数，然后把所有数往中位数靠近，维护大于，小于中位数的各有多少个和它们的值的和，然后大的值的和减去中位数乘以大的的个数，中位数乘以小的的个数减去小的值的和，两部分相加即为代价。
若某区间可实现在 $k$ 代价之内成为彩虹子数组，则其子区间区间也可实现（少操作一部分数，代价显然更小），所以可以用双指针来优化这一部分。
中位数部分，由于有增加也有删除，优先队列不太好做，我用的是可重集。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
#include<bits/stdc++.h>
#define fir first
#define sec second 
#define int long long
#define pii pair<int,int>
#define fep(i,s,e) for(int i=s;i<e;i++)
#define pef(i,s,e) for(int i=s;i>e;i--)
#define rep(i,s,e) for(int i=s;i<=e;i++)
#define per(i,s,e) for(int i=s;i>=e;i--)
namespace FastIO{
	template<typename T>inline void read(T &x){
		x=0;int f=1;char c=getchar();
		for(;!isdigit(c);c=getchar())if(c=='-')f=-1;
		for(;isdigit(c);c=getchar())x=(x<<1)+(x<<3)+(c^48);x*=f;
	}
	template<typename T,typename...Args>
    inline void read(T &x,Args&...args){
        read(x);
        read(args...);
    }
	template<typename T>void print(T x){
		if(x<0)x=-x,putchar('-');
		if(x>9)print(x/10);
		putchar((x%10)^48);
	}
}
using namespace std;
using namespace FastIO;
const int N=5e5+5;
int T,n,k,a[N],b[N];
struct node{
    multiset<int>lft,rit;
    int sml=0,smr=0;
    void add(int x){
        if(lft.empty()||x<=*lft.rbegin()){
            lft.insert(x);
            sml+=x;
        }
		else{
            rit.insert(x);
            smr+=x;
        }
        work();
    }
    void move(int x){
        if(lft.empty())return;
        int m=*lft.rbegin();
        if(x<=m){
            lft.erase(lft.find(x));
            sml-=x;
        }
		else{
            rit.erase(rit.find(x));
            smr-=x;
        }
        work();
    }
    void work(){
        if(lft.size()<rit.size()){
            auto it=rit.begin();
            int val=*it;
            lft.insert(val);
            sml+=val;
            smr-=val;
            rit.erase(it);
        }
		else if(lft.size()>rit.size()+1){
            auto it=prev(lft.end());
            int val=*it;
            rit.insert(val);
            smr+=val;
            sml-=val;
            lft.erase(it);
        }
    }
    int cost(){
        if(lft.empty())return 0;
        int m=*lft.rbegin();
        return m*lft.size()-sml+smr-m*rit.size();
    }
};
void solve(){
	read(n,k);
    fep(i,0,n){
    	read(a[i]);
    }
    fep(i,0,n){
        b[i]=a[i]-(i+1);
    }
    node mm;
    int l=0,ans=0;
    fep(r,0,n){
        mm.add(b[r]);
        while(mm.cost()>k&&l<=r){
            mm.move(b[l]);
            l++;
        }
        ans=max(ans,r-l+1);
    }
    print(ans);
    puts("");
}
signed main(){
    read(T);
    while(T--){
        solve();
    }
}