[ZJOI2019]开关

[ZJOI2019]开关

题面

洛谷

题解

令\(P=\sum p_i\)。

那么按到最终状态不停下继续往下按的 EGF 是

\[F(x)=\prod \frac {e^{\frac {p_i}Px}+e^{-\frac {p_i}Px}(-1)^{s_i}}2 \]

从某个状态按回自身状态的 EGF 为

\[G(x)=\prod \frac{e^{\frac {p_i}Px}+e^{-\frac {p_i}Px}}2 \]

我们把\(F(x)\)变为其 OGF \(f(x)\)（直接把\(x^i\)下面的\(i!\)）去掉，\(G(x)\)变为其 OGF \(g(x)\)，那么这分别就是\(F,G\)的概率生成函数。
答案的概率生成函数\(h(x)\)就是\(\frac fg\)，那么最终把\(h'(1)\)求出来就是答案。

然后考虑如何求上述的所有东西：

首先是 EGF 转OGF：

\[\begin{aligned} F(x)&=\sum_{i=-P}^Pa_ie^{\frac iPx}\\ &=\sum_{i=-P}^Pa_i\sum_{j\geq 0}\frac {(\frac {xi}P)^j}{j!} \end{aligned} \]

把\(j!\)去掉后直接等比数列求和：

\[f(x)=\sum_{i=-P}^P\frac {a_i}{1-\frac {xi}{P}} \]

然后是求\(h'(1)=(\frac {f(1)}{g(1)})'=\frac {f'(1)g(1)-f(1)g'(1)}{g(1)^2}\)。
考虑到当\(f(x),g(x),x=1\)时必有分母\(=1-\frac PP=0\)，所以我们将\(f,g\)同乘\(1-x\)，结果不影响而且当\(x=1\)时有意义。
因为\((\frac {1-x}{1-wx})'\large {|}_{x=1}=\frac {1}{w-1}\)
然后有\(f'(1)=\sum_{i=-P}^{P-1} \frac {a_i}{\frac iP-1}\)，\(g'(1)\)同理。
然后因为\(f,g\)均为概率生成函数所以\(f(1)=g(1)=1\)，所以我们就做完了（\(2^n\)也抵掉了）。

代码

#include <bits/stdc++.h> 
using namespace std; 
int gi() { 
	int res = 0, w = 1; 
	char ch = getchar(); 
	while (ch != '-' && !isdigit(ch)) ch = getchar(); 
	if (ch == '-') w = -1, ch = getchar(); 
	while (isdigit(ch)) res = res * 10 + ch - '0', ch = getchar(); 
	return res * w; 
} 
const int Mod = 998244353; 
int fpow(int x, int y) {
	int res = 1; 
	while (y) {
		if (y & 1) res = 1ll * res * x % Mod; 
		x = 1ll * x * x % Mod, y >>= 1; 
	} 
	return res; 
} 
const int MAX_N = 1e5 + 5; 
void Pls(int &x, int y) { x += y; if (x >= Mod) x -= Mod; } 
int N, s[MAX_N], p[MAX_N], P; 
int f[MAX_N], g[MAX_N], tf[MAX_N], tg[MAX_N]; 
int main () { 
#ifndef ONLINE_JUDGE 
    freopen("cpp.in", "r", stdin); 
#endif 
	N = gi(); 
	for (int i = 1; i <= N; i++) s[i] = gi(); 
	for (int i = 1; i <= N; i++) p[i] = gi(), P += p[i]; 
	f[P] = g[P] = 1; 
	for (int i = 1, k = 0; i <= N; i++) { 
		k += p[i]; 
		for (int j = 0; j <= P << 1; j++) tf[j] = tg[j] = 0; 
		for (int j = 0; j <= P << 1; j++) { 
			if (j - p[i] >= 0) Pls(tf[j], f[j - p[i]]); 
			if (j + p[i] <= P << 1) Pls(tf[j], s[i] ? (-f[j + p[i]] + Mod) : f[j + p[i]]);  
		} 
		for (int j = 0; j <= P << 1; j++) { 
			if (j - p[i] >= 0) Pls(tg[j], g[j - p[i]]); 
			if (j + p[i] <= P << 1) Pls(tg[j], g[j + p[i]]);  
		} 
		for (int j = 0; j <= P << 1; j++) f[j] = tf[j], g[j] = tg[j]; 
	} 
	int ans = 0, inv = fpow(P, Mod - 2); 
	for (int i = -P; i < P; i++) Pls(ans, 1ll * f[i + P] * fpow(1ll * (i + Mod) * inv % Mod - 1, Mod - 2) % Mod); 
	for (int i = -P; i < P; i++) Pls(ans, Mod - 1ll * g[i + P] * fpow(1ll * (i + Mod) * inv % Mod - 1, Mod - 2) % Mod); 
	printf("%d\n", ans); 
    return 0; 
}