题解：AT_arc139_d [ARC139D] Priority Queue 2

题面

---

发现我们不好算到最后还剩些什么。

考虑计算 \(\sum\limits_{i=1}^m\sum\limits_{j=1}^n[s_j\ge i]\)，容易发现这和原式等价。

记 \(b_i\) 表示 \(s\) 中不小于 \(i\) 的数的个数，每次删去第 \(x\) 小的等价于将所有超过 \(n-x+1\) 的地方减 1，加入 \(k\) 等价于将 \(b_{1,k}\) 都加 1。

发现最终的答案就是 \(\sum b\)。

考虑计算每一个位值的贡献，以下记 \(x'=n-x+1\)。

枚举每个位值被加了多少次，倘若最初 \(b_i\le x'\)，有 \(b_i'=\min(b_i+j,x')\)；如果最初 \(b_i>x'\)，有 \(b_i'=\max(x,b_i+j-k)\)。其中 \(j\) 是第 \(i\) 个位值被加的次数，\(k\) 是总操作次数。

---

代码：

#include<bits/stdc++.h>
#define int long long
#define pii pair<int,int>
#define mp(a,b) make_pair(a,b)
#define db long double
using namespace std;
const int mod=998244353;
int n,m,k,x,a[2005];
int c[2005],inc[2005];
int fpow(int a,int b){
	int res=1;
	while(b){
		if(b&1)res=res*a%mod;
		a=a*a%mod;
		b>>=1;
	}
	return res;
}
int C(int n,int m){
	return c[n]*inc[n-m]%mod*inc[m]%mod;
}
signed main()
{
	c[0]=1;
	for(int i=1;i<=2000;i++)c[i]=c[i-1]*i%mod;
	inc[2000]=fpow(c[2000],mod-2);
	for(int i=1999;i>=0;i--)inc[i]=inc[i+1]*(i+1)%mod;
	scanf("%lld%lld%lld%lld",&n,&m,&k,&x);
	x=n+1-x;
	for(int i=1,c;i<=n;i++)scanf("%lld",&c),a[c]++;
	for(int i=m;i>=1;i--)a[i]+=a[i+1];
	int ans=0;
	for(int i=1;i<=m;i++){
		int l=m-i+1,r=i-1;
		for(int j=0;j<=k;j++){
			if(a[i]<=x)ans=(ans+(min(x,a[i]+j)*C(k,j)%mod*fpow(l,j)%mod*fpow(r,k-j)%mod))%mod;
			else ans=(ans+(max(x,a[i]+j-k)*C(k,j)%mod*fpow(l,j)%mod*fpow(r,k-j)%mod))%mod;
		}
	}
	printf("%lld",ans);
	return 0;
}