[Luogu4705] 玩游戏

Description

给定两个长度分别为 \(n\) 和 \(m\) 的序列 \(a\) 和 \(b\)。要从这两个序列中分别随机一个数，设为 \(a_x,b_y\)，定义该次游戏的 \(k\) 次收益为 \((a_x+b_y)^k\) 。对于 \(i=1,2,\dots,t\)，求一次游戏 \(i\) 次收益的期望。\(n,m,t\leq 10^5\)。

Sol

根据期望的线性性，显然可以求每个点对的 \(i\) 次收益，最后再除以 \(nm\) 就好了。

所以问题转化为，对于每个 \(k\)，求：

\[\sum_{i=1}^n\sum_{j=1}^m (a_i+b_j)^k \]

接下来直接推导：

\[\begin{aligned} ans_k&=\sum_{i=1}^n\sum_{j=1}^m (a_i+b_j)^k\\ &=\sum_{i=1}^n\sum_{j=1}^m\sum_{p=0}^k \binom kpa_i^pb_j^{k-p}\\ &=\sum_{p=0}^k\binom kp \left(\sum_{i=1}^na_i^p\right) \left(\sum_{j=1}^mb_j^{k-p} \right)\\ &=k!\cdot\sum_{p=0}^k \left(\sum_{i=1} ^n \frac{a_i^p}{p!}\right) \left(\sum_{j=1}^m\frac{b_j^{k-p}}{(k-p)!} \right) \end{aligned} \]

发现这是个卷积式子，现在问题变成了如何求：

\[\sum_{i=1}^n a_i^p \]

设 \(F(x)=\prod\limits_{i=1}^n(1+a_ix),G(x)=\ln(F(x))\)

那么：

\[\begin{aligned} G(x)&=\ln(\prod_{i=1}^n 1+a_ix)\\ &=\sum_{i=1}^n \ln(1+a_ix) \end{aligned} \]

把 \(\ln(1+a_ix)\) 泰勒展开：

\[\begin{aligned} G(x)&=\sum_{i=1}^n \ln(1+a_ix)\\ &= \sum_{i=1}^n \sum_{k=1}^\infty \frac{(-1)^{k+1}}{k}\cdot a_i^k\cdot x^k\\ &= \sum_{k=1}^\infty \frac{(-1)^{k+1}}k\cdot x^k\cdot \left( \sum_{i=1}^n a_i^k \right) \end{aligned} \]

后边那项就是我们要求的。

总结一下，先分治\(\text{NTT}\)求出\(F(x)\)，再取对数求出\(G(x)\)，然后第 \(k\) 项乘上一个系数就是 \(\sum\limits_{i=1}^n a_i^k\) 了。

Code

#pragma GCC optimize(2)
#include<bits/stdc++.h>
using namespace std;
typedef double db;
typedef long long ll;
typedef vector<int> vec;
const int N=262144+5;
const int mod=998244353;
#define pb push_back

int w[2][N],in[N];
int fac[N],ifac[N],A[N],B[N];
int n,m,t,a[N],b[N],c[N],d[N];
int lim,maxn,rev[N],tmpa[N],tmpb[N];

int ksm(int a,int b=mod-2,int ans=1){
    while(b){
        if(b&1) ans=1ll*ans*a%mod;
        a=1ll*a*a%mod;b>>=1;
    } return ans;
}

void ntt(int *f,int g){
    for(int i=1;i<lim;i++) if(i<rev[i]) swap(f[i],f[rev[i]]);
    for(int mid=1;mid<lim;mid<<=1){
        for(int R=mid<<1,j=0;j<lim;j+=R){
            for(int k=0;k<mid;k++){
                int x=f[j+k],y=1ll*w[g][maxn/R*k]*f[j+k+mid]%mod;
                f[j+k]=x+y>=mod?x+y-mod:x+y,f[j+k+mid]=x-y<0?x-y+mod:x-y;
            }
        }
    } if(g)
        for(int i=0;i<lim;i++) f[i]=1ll*f[i]*in[lim]%mod;
}

vec calc(int *a,int l,int r){
    if(l==r){vec now;now.pb(1);now.pb(a[l]);return now;}
    int mid=l+r>>1;
    vec L=calc(a,l,mid),R=calc(a,mid+1,r);
    lim=1;while(lim<=r-l+1) lim<<=1;
    for(int i=1;i<lim;i++) rev[i]=(rev[i>>1]>>1)|(i&1?lim>>1:0);
    for(int i=0;i<(int)L.size();i++) A[i]=L[i];
    for(int i=0;i<(int)R.size();i++) B[i]=R[i];
    ntt(A,0),ntt(B,0);
    for(int i=0;i<lim;i++) A[i]=1ll*A[i]*B[i]%mod;
    ntt(A,1); vec now;
    for(int i=0;i<=r-l+1;i++) now.pb(A[i]),A[i]=B[i]=0;
    for(int i=r-l+2;i<lim;i++) A[i]=B[i]=0;
    return now;
}

void solveinv(int *a,int *b,int len){
    if(len==1) return b[0]=ksm(a[0]),void();
    solveinv(a,b,len>>1); lim=len<<1;
    for(int i=1;i<lim;i++) rev[i]=(rev[i>>1]>>1)|(i&1?lim>>1:0);
    for(int i=len;i<lim;i++) A[i]=0;
    for(int i=0;i<len;i++) A[i]=a[i];
    ntt(A,0),ntt(b,0);
    for(int i=0;i<lim;i++)
        b[i]=1ll*b[i]*(2ll-1ll*A[i]*b[i]%mod+mod)%mod;
    ntt(b,1); for(int i=len;i<lim;i++) b[i]=0;
}

void ds(int *a,int *b,int n){
    for(int i=0;i<n;i++)
        b[i]=1ll*a[i+1]*(i+1)%mod;
    b[n]=0;
}

void jf(int *a,int n){
    for(int i=n;i;i--)
        a[i]=1ll*a[i-1]*in[i]%mod;
    a[0]=0;
}

void solveln(int *a,int *b,int n){
    memset(tmpa,0,sizeof tmpa);
    memset(tmpb,0,sizeof tmpb);
    lim=1;while(lim<n) lim<<=1;
    solveinv(a,tmpa,lim);
    lim=1;while(lim<n<<1) lim<<=1;
    for(int i=1;i<lim;i++) rev[i]=(rev[i>>1]>>1)|(i&1?lim>>1:0);
    ds(a,tmpb,n);
    ntt(tmpa,0),ntt(tmpb,0);
    for(int i=0;i<lim;i++) b[i]=1ll*tmpa[i]*tmpb[i]%mod;
    ntt(b,1); jf(b,n);
}

void init(int n){
    fac[0]=ifac[0]=1;
    for(int i=1;i<=n;i++) fac[i]=1ll*fac[i-1]*i%mod;
    ifac[n]=ksm(fac[n]);
    for(int i=n-1;i;i--) ifac[i]=1ll*ifac[i+1]*(i+1)%mod;
}

signed main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++) scanf("%d",&a[i]);
    for(int i=1;i<=m;i++) scanf("%d",&b[i]);
    scanf("%d",&t);
    init(t);
    maxn=1;while(maxn<=max(t<<1,n+m-2)) maxn<<=1;
    w[0][0]=w[1][0]=1; in[1]=1;
    w[0][1]=ksm(3,(mod-1)/maxn),w[1][1]=ksm((mod+1)/3,(mod-1)/maxn);
    for(int i=2;i<=maxn;i++)
    	in[i]=ksm(i),
        w[0][i]=1ll*w[0][i-1]*w[0][1]%mod,
        w[1][i]=1ll*w[1][i-1]*w[1][1]%mod;
    vec aa=calc(a,1,n),bb=calc(b,1,m);
    for(int i=0;i<=n;i++) c[i]=aa[i];
    for(int i=0;i<=m;i++) d[i]=bb[i];
    memset(a,0,sizeof a),memset(b,0,sizeof b);
    solveln(c,a,t); a[0]=n; // 注意这里的0次项 积分给消掉了 所以要特殊赋值
    solveln(d,b,t); b[0]=m; 
    for(int i=1;i<=t;i++){
        a[i]=1ll*a[i]*i%mod;
        b[i]=1ll*b[i]*i%mod;
        if(!(i&1)) a[i]=mod-a[i],b[i]=mod-b[i];
        a[i]=1ll*a[i]*ifac[i]%mod;
        b[i]=1ll*b[i]*ifac[i]%mod;
    }
    for(int i=t+1;i<lim;i++) a[i]=b[i]=0;
    lim=maxn;
    for(int i=1;i<lim;i++) rev[i]=(rev[i>>1]>>1)|(i&1?lim>>1:0);
    ntt(a,0),ntt(b,0);
    for(int i=0;i<lim;i++) a[i]=1ll*a[i]*b[i]%mod;
    ntt(a,1);
    for(int inn=ksm(1ll*n*m%mod),i=1;i<=t;i++) 
        printf("%lld\n",1ll*a[i]*fac[i]%mod*inn%mod);
    return 0;
}