题解:AT_arc139_d [ARC139D] Priority Queue 2

· · 题解

直接从集合的角度考虑计算总和不好处理。考虑换一个方向,因为元素的值域很小,所以可以从每种数的角度计算。具体来说,我们求出 s_i 表示所有可能的 A 中值等于 i 的位置的个数之和,然后计算答案使用 \sum is_i

套路地,把等于转化为大于等于。我们枚举 k,需要计算所有 A\ge k 的元素个数之和。先对一个单独的集合考虑,记现在 \ge k 个数为 c,那么有两种情况:加入了 \ge k 的数,c 加一;否则,c 不变。操作完后,我们需要删除第 X 小的数,那么如果 c>n+1−Xc 就会减一。

考虑这种操作方式的性质:如果给出一个集合初始 \ge k 的元素个数,以及加入 \ge k 的数的次数,最后的 c 是固定的,与什么时候加入 \ge k 的数无关。于是可以再枚举加入 \ge k 数的次数 t,此时可以计算出最后的 c,而这种情况的总方案数为 (M−i+1)^t(i−1)^{K−t}\binom{K}{t}。于是最终的答案就可以算了。

AC code:

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int mo=998244353;
int n,m,k,x,a[5001],cn;
int an,sc[5001],iv[5001];
void pr(){
    for(int i=2;i<=k;i++)
        iv[i]=(mo-(mo/i)*iv[mo%i]%mo)%mo;
    for(int i=1;i<=k;i++)
        sc[i]=sc[i-1]*(k-i+1)%mo*iv[i]%mo;
}
int pw(int a,int b){
    int re=1;
    while(b){
        if(b&1)re=re*a%mo;
        a=a*a%mo,b>>=1;
    }
    return re;
}
signed main(){
    ios::sync_with_stdio(0);
    cin>>n>>m>>k>>x,iv[1]=sc[0]=1,pr();
    for(int i=1;i<=n;i++)cin>>a[i];
    for(int i=1;i<=m;i++,cn=0){
        for(int j=1;j<=n;j++)cn+=(a[j]>=i);
        for(int j=0;j<=k;j++){
            int c=max(cn+j-k,min(cn+j,n+1-x));
            an=(an+c*pw(i-1,k-j)%mo
                 *pw(m-i+1,j)%mo*sc[j]%mo)%mo;
        }
    }
    cout<<an;
    return 0;
}