ABC471Eの题解

题目传送门

简要题意

给定 \(N\) 个数，在所有恰好选 \(K\) 个元素的子集里，求出每个子集元素和的平方之和，答案对 \(998244353\) 取模。

做法

首先我们暴力统计是会超时的，所以我们考虑拆贡献，先将平方拆开，得到 \(\sum_i A_i^2+2\times\sum_{i,j}A_i\times A_j\)
选到 \(A_i^2\) 的总贡献是 \(\binom{n-1}{k-1}\)
选到 \(2\times A_i\times A_j\) 的总贡献是 \(\binom{n-2}{k-2}\)

代码

//由AI生成
#include<bits/stdc++.h>
using namespace std;
using ll=long long;
const int MOD=998244353;
const int MAX=200010;

ll fac[MAX],invfac[MAX];

ll powmod(ll a,ll b){
    ll res=1;
    while(b>0){
        if(b&1) res=res*a%MOD;
        a=a*a%MOD;
        b>>=1;
    }
    return res;
}
void pre(int n){
    fac[0]=1;
    for(int i=1;i<=n;i++) fac[i]=fac[i-1]*i%MOD;
    invfac[n]=powmod(fac[n],MOD-2);
    for(int i=n-1;i>=0;i--) invfac[i]=invfac[i+1]*(i+1)%MOD;
}
ll C(int n,int m){
    if(n<0||m<0||n<m) return 0;
    return fac[n]*invfac[m]%MOD*invfac[n-m]%MOD;
}
int main(){
    ios::sync_with_stdio(false);cin.tie(0);
    int N,K;cin>>N>>K;
    pre(N);
    ll S=0,s2=0;
    for(int i=1;i<=N;i++){
        ll a;cin>>a;
        a%=MOD;
        S=(S+a)%MOD;
        s2=(s2+a*a%MOD)%MOD;
    }
    ll c1=C(N-1,K-1);
    ll c2=C(N-2,K-2);
    ll part1=c1*s2%MOD;
    ll cross=(S*S%MOD - s2 + MOD)%MOD;
    ll part2=c2*cross%MOD;
    ll ans=(part1+part2)%MOD;
    cout<<ans<<endl;
    return 0;
}