【LibreOJ #3045.「ZJOI2019」开关】【生成函数】

题意

有$n$个初始均处于关闭状态的开关，每个开关有一个权值$p_i$，每一轮会随机选择一个开关并改变其状态，且抽中第$i$个开关的概率为$\frac{p_i}{\sum p_i}$。问变为目标状态期望需要经过多少轮。
$n\le 100,\sum p_i\le 5\ast 10^4$

分析

设$F(x)$表示经过$n$轮后恰好变为目标状态的概率的$EGF$，$G(x)$表示经过$n$轮后恰好回到初始状态的概率的$EGF$，目标状态为$t_i$，显然有$$F(x)=\prod _{i=1}^n\frac{e^{p_{i}x}+(-1{)}^{t_i}{e}^{-p_{i}x}}{2}$$
$$G(x)=\prod _{i=1}^n\frac{e^{p_{i}x}+e^{-p_{i}x}}{2}$$
设$H(x)$为经过$n$轮后第一次变为目标状态的概率的$EGF$，$f(x),g(x),h(x)$分别表示$F(x),G(x),H(x)$的$OGF$，那么有$$f(x)=g(x)h(x)$$
从而得到$$h(x)=\frac{f(x)}{g(x)}$$
显然我们要求的就是$h^{\prime }(1)$
考虑如果知道$EGF$怎么求$OGF$，设$F(x)=\sum a{e}^{vx}$，则$f(x)=\sum \frac{a}{1-vx}$
于是我们可以先通过dp求出$F(x)$和$G(x)$，那么$$h^{\prime }=(\frac{f}{g}{)}^{\prime }=\frac{f^{\prime }g-f{g}^{\prime }}{g^2}$$
注意到$f(x)$和$g(x)$的分母中可能含有$(1-x)$的项，可以先把两边同乘$(1-x)$，通过推导容易得到$$(\frac{1-x}{1-vx}{)}^{\prime }{|}_{x=1}=\frac{1}{v-1}$$
时间复杂度$O(n\sum p_i)$

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>

typedef long long LL;

const int N=50005;
const int MOD=998244353;

int n,sum,F[N*2],G[N*2],t[N],p[N],tmp[N*2];

int ksm(int x,int y)
{
	int ans=1;
	while (y)
	{
		if (y&1) ans=(LL)ans*x%MOD;
		x=(LL)x*x%MOD;y>>=1;
	}
	return ans;
}

void solve()
{
	int inv=ksm(2,MOD-2);
	F[sum]=1;
	for (int i=1;i<=n;i++)
	{
		int x=inv,y=!t[i]?inv:MOD-inv;
		for (int j=0;j<=sum*2;j++)
		{
			tmp[j]=0;
			if (j>=p[i]) (tmp[j]+=(LL)x*F[j-p[i]]%MOD)%=MOD;
			if (j+p[i]<=sum*2) (tmp[j]+=(LL)y*F[j+p[i]]%MOD)%=MOD;
		}
		for (int j=0;j<=sum*2;j++) F[j]=tmp[j];
	}
	G[sum]=1;
	for (int i=1;i<=n;i++)
	{
		int x=inv,y=inv;
		for (int j=0;j<=sum*2;j++)
		{
			tmp[j]=0;
			if (j>=p[i]) (tmp[j]+=(LL)x*G[j-p[i]]%MOD)%=MOD;
			if (j+p[i]<=sum*2) (tmp[j]+=(LL)y*G[j+p[i]]%MOD)%=MOD;
		}
		for (int j=0;j<=sum*2;j++) G[j]=tmp[j];
	}
}

int calc()
{
	int inv=ksm(sum,MOD-2),f1=0,g1=0,f2=0,g2=0;
	for (int i=0;i<=sum*2;i++)
	{
		int x=(LL)(i-sum)*inv%MOD;
		x+=x<0?MOD:0;
		if (x==1) (f1+=F[i])%=MOD,(g1+=G[i])%=MOD;
		else (f2+=(LL)F[i]*ksm(x-1,MOD-2)%MOD)%=MOD,(g2+=(LL)G[i]*ksm(x-1,MOD-2)%MOD)%=MOD;
	}
	return (LL)((LL)f2*g1%MOD-(LL)f1*g2%MOD)*ksm((LL)g1*g1%MOD,MOD-2)%MOD;
}

int main()
{
	scanf("%d",&n);
	for (int i=1;i<=n;i++) scanf("%d",&t[i]);
	for (int i=1;i<=n;i++) scanf("%d",&p[i]),sum+=p[i];
	solve();
	printf("%d\n",(calc()+MOD)%MOD);
	return 0;
}