[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;
}