[ABC231G] Balls in Boxes

· · 题解

[ABC231G] Balls in Boxes

题解

我们记 b_i 为 i 被选中的次数，那么期望得分为

\frac{1}{n^k}\sum\limits_{b_1+b_2+\dots+b_n=k}\dbinom{n}{b_1,b_2,\dots,b_n}\prod\limits_{i=1}^n(a_i+b_i) =\frac{k!}{n^k}\sum\limits_{b_1,b_2,\dots,b_n}\prod\limits_{i=1}^n\frac{a_i+b_i}{b_i!}

考虑构造 EGF：

\hat F_i(x)=\sum\limits_{j=0}^{\infin}\frac{a_i+j}{j!}x^j =a_i\sum\limits_{j=0}^{\infin}\frac{1}{j!}x^j+\sum\limits_{j=0}^{\infin}\frac{1}{(j-1)!}x^j =a_ie^x+xe^x =(a_i+x)e^x

那么答案转化为：

[x^k]\left(\frac{k!}{n^k}\prod\limits_{i=1}^n\hat F_i(x)\right) =[x^k]\left(\frac{k!}{n^k}e^{nx}\prod\limits_{i=1}^n(a_i+x)\right)

我们把 \prod\limits_{i=1}^n(a_i+x) 展开成 \prod\limits_{i=0}^{n}c_ix^i。这一步可以 O(n^2) 暴力做，也可以用一些多项式做法做到 O(n\log^2n)。

再把 e 展开，答案为：

=\frac{k!}{n^k}\sum\limits_{i=0}^{\min(n,k)}c_i\frac{n^{k-i}}{(k-i)!} =\sum\limits_{i=0}^{\min(n,k)}c_ik^{\underline{i}}n^{-i}

代码

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=1010,mod=998244353;
ll n,k,c[N],a[N];
ll qpow(ll x,ll y){
    ll ret=1;
    while(y){
        if(y&1)ret=ret*x%mod;
        x=x*x%mod;
        y>>=1; 
    }
    return ret;
}
int main(){
    cin>>n>>k;
    c[0]=1;
    for(int i=1;i<=n;i++){
        cin>>a[i];
        for(int j=n;j;j--)c[j]=(a[i]*c[j]+c[j-1])%mod;
        c[0]=c[0]*a[i]%mod;
    }
    ll ans=0,fac=1,p=1,inv=qpow(n,mod-2);
    for(int i=0;i<=min(n,k);i++){
        ans=(ans+c[i]*fac%mod*p%mod)%mod;
        fac=fac*(k-i)%mod;p=p*inv%mod;
    }
    cout<<ans;
}