[ZJOI2019]开关

Description

有一个n位二进制数，一开始为0
每一轮中对于每一个i，有$\frac{p_i}{\sum p_i}$的概率将这一位异或上1，每轮操作独立
问第一次变成目标状态的期望步数
n<=100,∑pi<=50000

Solution

考虑设F(x)表示在n步时变为目标状态的概率，显然有$$F(x)=\prod \frac{e^{p_{i}x}+(-1{)}^{t_i}{e}^{-p_{i}x}}{2}$$这里的pi都是概率
这样不能保证是第一次，所以我们再设G(x)表示n步内变回0状态的概率，那么$$G(x)=\prod \frac{e^{p_{i}x}+e^{-p_{i}x}}{2}$$
设答案的EGF为H(x)，其对应的OGF $f(x),g(x),h(x)$,会满足$h\ast g=f->h=f/g$
假如我们求出了答案的ogf h(x)，那么问题的答案就是$h^{\prime }(1)$
即我们对右边求导，将x=1代入
考虑知道EGF怎么求OGF，若$F(x)=\sum a{e}^{vx}$那么显然有$f(x)=\sum \frac{a}{1-vx}$
那么我们直接Dp求出EGF的表达式自然就知道了OGF的表达式
考虑$(\frac{f(x)}{g(x)}{)}^{\prime }=\frac{f^{\prime }(x)g(x)-f(x)g^{\prime }(x)}{g^2(x)}$，我们只需要知道$f(1),f^{\prime }(1),g(1),g^{\prime }(1)$就能算
但问题是$\frac{a}{1-vx}$可能在1处未定义，怎么办呢？
我们可以直接通分，令$f(x)=\frac{A(x)}{B(x)},g(x)=\frac{C(x)}{D(x)}$，注意到B(x)=D(x)，所以我们等价于要算$\frac{(A(x)}{C(x){)}^{\prime }}$，也就是A(1),A’(1),C(1),C’(1)，这些直接维护一个1-v的前缀积就好了
复杂度O(n∑pi)，当然用NTT也可以做到O(∑pi log^2 ∑pi)

Code

#include <cstdio>
#include <cstring>
#include <algorithm>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
using namespace std;

typedef long long ll;

const int N=1e5+5,Mo=998244353;

int pwr(int x,int y) {
	int z=1;
	for(;y;y>>=1,x=(ll)x*x%Mo)
		if (y&1) z=(ll)z*x%Mo;
	return z;
}

int n,m,w,sum,f[N],g[N],t[N],p[N],pre[2][N],suf[2][N],v[N],A,C,dA,dC;

int main() {
	scanf("%d",&n);
	fo(i,1,n) scanf("%d",&t[i]);
	fo(i,1,n) scanf("%d",&p[i]),sum+=p[i];
	f[0]=g[0]=1;m=0;
	fo(i,1,n) {
		m+=p[i]<<=1;
		fd(j,m,p[i]) f[j]=(f[j-p[i]]+((t[i]&1)?-f[j]:f[j]))%Mo;
		fd(j,p[i]-1,0) f[j]=(t[i]&1)?-f[j]:f[j];
		fd(j,m,p[i]) g[j]=(g[j-p[i]]+g[j])%Mo;
	}
	fo(i,0,m) 
		if (!(i&1)) {
			v[i/2]=(ll)(i-m/2)*pwr(sum,Mo-2)%Mo;
			f[i/2]=f[i];g[i/2]=g[i];
		}
	m/=2;
	pre[0][0]=1-v[0];pre[1][0]=-v[0];
	fo(i,1,m) {
		pre[0][i]=(ll)pre[0][i-1]*(1-v[i])%Mo;
		pre[1][i]=(ll)pre[1][i-1]*(1-v[i])%Mo;
		(pre[1][i]-=(ll)pre[0][i-1]*v[i]%Mo)%=Mo;
	}
	suf[0][m]=1-v[m];suf[1][m]=-v[m];
	fd(i,m-1,0) {
		suf[0][i]=(ll)suf[0][i+1]*(1-v[i])%Mo;
		suf[1][i]=(ll)suf[1][i+1]*(1-v[i])%Mo;
		(suf[1][i]-=(ll)suf[0][i+1]*v[i]%Mo)%=Mo;
	}
	fo(i,0,m) {
		int now=(ll)(i?pre[0][i-1]:1)*(i<m?suf[0][i+1]:1)%Mo;
		(A+=(ll)f[i]*now%Mo)%=Mo;(C+=(ll)g[i]*now%Mo)%=Mo;
		now=(ll)(i?pre[0][i-1]:1)*(i<m?suf[1][i+1]:0)%Mo;
		(now+=(ll)(i?pre[1][i-1]:0)*(i<m?suf[0][i+1]:1)%Mo)%=Mo;
		(dA+=(ll)f[i]*now%Mo)%=Mo;(dC+=(ll)g[i]*now%Mo)%=Mo;
	}
	int ans=((ll)dA*C%Mo-(ll)A*dC%Mo)%Mo;
	ans=(ll)ans*pwr(pwr(C,2),Mo-2)%Mo;
	printf("%d\n",(ans+Mo)%Mo);
	return 0;
}