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

传送门

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

• 题解：
  我们令成功了但不结束的概率指数生成函数为 $F(x)$，直接上单位根反演
  $F(x)=\prod_{i=1}^n\sum_{j=0}^{\infty}\frac{1+(-1)^{s_i+j}}{2}*(\frac{p_i}{sum})^j*\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)*g(x)\Rightarrow h(x)=\frac{f(x)}{g(x)}$
  注意到答案求的就是
  $h'(1)=\frac{-g'(1)f(1)+f'(1)g(1)}{g(1)^2}$
  下面考虑分别求出 $f(x),g(x),f'(x),g'(x)$
  当 $F(x)=\sum_{i=-sum}^{sum}a_i*e^{\frac{i}{sum}x}$ 时
  $f(x)=\sum_{i=-sum}^{sum}a_i*\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\neq i} (1-\frac{j}{sum}x)$
  只在 $i=sum$ 处有值
  $f'(x)=\sum_{i=-sum}^{sum}a_i\sum_{j\neq i}-\frac{j}{sum}\prod_{k\neq i,k\neq 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;
}