「ZJOI2019」开关 （生成函数）（概率）

传送门

一道毒瘤的生成函数推导题
一开始发现概率 $dp$ 加一个高斯消元可以过 $50pts$，然后一直在想 $dp$……

• 题解：
  我们令成功了但不结束的概率指数生成函数为 $F(x)$，直接上单位根反演
  $$F(x)=\prod _{i=1}^n\sum _{j=0}^{\infty }\frac{1+(-1{)}^{s_i+j}}{2}\ast (\frac{p_i}{sum}{)}^j\ast \frac{x^j}{j!}$$
  $$=\prod _{i=1}^n\frac{e^{\frac{p_i}{sum}x}+(-1{)}^{s_i}{e}^{-\frac{p_i}{sum}x}}{2}$$
  考虑让成功过后停下来，我们构造走 $i$ 步回到本身状态的概率指数生成函数 $G(x)$，那么
  $$G(x)=\prod _{i=1}^n\frac{e^{\frac{p_i}{sum}x}+e^{-\frac{p_i}{sum}x}}{2}$$
  令走 $i$ 步成功的概率的普通生成从函数为 $h(x)$，$F(x),G(x)$转为普通生成函数为 $f(x),g(x)$，
  那么有
  $$f(x)=h(x)\ast g(x)\Rightarrow h(x)=\frac{f(x)}{g(x)}$$
  注意到答案求的就是
  $$h^{\prime }(1)=\frac{-g^{\prime }(1)f(1)+f^{\prime }(1)g(1)}{g(1{)}^2}$$
  下面考虑分别求出 $f(x),g(x),f^{\prime }(x),g^{\prime }(x)$
  当 $F(x)={\sum }_{i=-sum}^{sum}{a}_i\ast e^{\frac{i}{sum}x}$ 时
  $$f(x)=\sum _{i=-sum}^{sum}{a}_i\ast \frac{1}{1-\frac{i}{sum}x}$$
  考虑这玩意在 1 处不收敛，将 $f,g$ 同乘 $\prod (1-\frac{i}{sum}x)$
  那么
  $$f(x)=\sum _{i=-sum}^{sum}{a}_i\prod _{j\ne i}(1-\frac{j}{sum}x)$$
  只在 $i=sum$ 处有值
  $$f^{\prime }(x)=\sum _{i=-sum}^{sum}{a}_i\sum _{j\ne i}-\frac{j}{sum}\prod _{k\ne i,k\ne j}(1-\frac{k}{sum}x)$$
  只在 $i=sum$ 或 $j=sum$ 处有值，于是可以 $O(\sum p_i)$ 做
  处理系数 $a_i$ 就是个背包，复杂度 $O(n\sum p_i)$

#include<bits/stdc++.h>
#define cs const
using namespace std;
int read(){
	int cnt = 0, f = 1; char ch = 0;
	while(!isdigit(ch)){ ch = getchar(); if(ch == '-') f = -1; }
	while(isdigit(ch)) cnt = cnt*10 + (ch-'0'), ch = getchar();
	return cnt * f;
}
cs int N = 105, M = 1e5 + 50;
cs int Mod = 998244353;
int add(int a, int b){ return a + b >= Mod ? a + b - Mod : a + b; }
int mul(int a, int b){ return 1ll * a * b % Mod; }
int dec(int a, int b){ return a - b < 0 ? a - b + Mod : a - b; }
int ksm(int a, int b){ int ans=1; for(;b;b>>=1,a=mul(a,a)) if(b&1) ans=mul(ans,a); return ans; }
void Add(int &a, int b){ a = add(a, b); }
void Dec(int &a, int b){ a = dec(a, b); }
void Mul(int &a, int b){ a = mul(a, b); }
int fix(int a){ return a < 0 ? a + Mod : a; }
int n, s[N], p[N], S, iS, f[M], g[M];
int func(int *f){
	int ans = f[S+S];
	for(int i = -S; i < S; i++) Mul(ans, dec(1,mul(fix(i),iS)));
	return ans;
}
int deriv(int *f){
	static int coef[M];
	int sum = 1, ans = 0;
	for(int i = -S; i < S; i++) Mul(sum, dec(1,mul(fix(i),iS)));
	for(int i = -S; i < S; i++) coef[i+S] = mul(sum, ksm(dec(1,mul(fix(i),iS)),Mod-2));
	sum = 0;
	for(int i = -S; i < S; i++) Add(sum, mul(mul(fix(i),iS), coef[i+S]));
	Add(ans, mul(sum, f[S+S]));
	for(int i = -S; i < S; i++) Dec(ans, mul(f[i+S], coef[i+S]));
	return ans;
}
int main(){
	n = read();
	for(int i = 1; i <= n; i++) s[i] = read();
	for(int i = 1; i <= n; i++) p[i] = read(), S += p[i];
	iS = ksm(S, Mod-2);
	f[S] = g[S] = 1;
	for(int i = 1, now = 0; i <= n; i++){
		static int tf[M], tg[M];
		memset(tf, 0, sizeof(tf)); 
		memset(tg, 0, sizeof(tg));
		for(int j = -now; j <= now; j++){
			if(s[i] == 0) Add(tf[j + S - p[i]], f[j + S]);
			else Dec(tf[j + S - p[i]], f[j + S]);
			Add(tf[j + S + p[i]], f[j + S]);
			Add(tg[j + S + p[i]], g[j + S]);
			Add(tg[j + S - p[i]], g[j + S]);
		}
		memcpy(f, tf, sizeof(tf));
		memcpy(g, tg, sizeof(tg)); now += p[i];
	}
	int fg = func(g), res = mul(fg, deriv(f));
	Dec(res, mul(func(f), deriv(g)));
	cout << mul(res, ksm(mul(fg,fg), Mod-2)); return 0;
}