【LOJ3045】「ZJOI2019」开关

【题目链接】

• 点击打开链接

【思路要点】

• 无视结束条件，令 $i$ 步之后状态全零的概率的指数型生成函数为 $F(x)$ ，对应的普通型生成函数为 $f(x)$ ，记 $S=\sum p_i$ ，则有 $$F(x)=\prod _{i=1}^N\frac{exp(\frac{p_i}{S}x)+(-1{)}^{s_i}exp(-\frac{p_i}{S}x)}{2}$$
• 考虑去掉不是首次达到状态全零的概率，记 $i$ 步之后状态不变的概率的指数型生成函数为 $G(x)$ ，首次达到状态全零的概率的指数型生成函数为 $H(x)$ ，对应的普通型生成函数分别为 $g(x),h(x)$ 。显然，有 $$G(x)=\prod _{i=1}^N\frac{exp(\frac{p_i}{S}x)+exp(-\frac{p_i}{S}x)}{2}$$
  并且 $$h(x)g(x)=f(x),h(x)=\frac{f(x)}{g(x)}$$
• 需要计算的期望即为 $h^{\prime }(1)$ 的值，即考虑计算 $$\frac{f^{\prime }(1)g(1)-f(1)g^{\prime }(1)}{g^2(1)}$$
• $F(x),G(x)$ 可以用定义式直接写作 ${\sum }_{i=-S}^S{a}_{i}exp(\frac{i}{S}x)$ ，考虑 $exp$ 的组合意义，将指数型生成函数转换为普通型生成函数，得到的结果应为 $$\sum _{i=-S}^S\frac{a_i}{1-\frac{i}{S}x}$$
• 直接将 $x=1$ 代入 $f(x),g(x)$ 求得的值趋于无穷，但注意到 $h(x)=\frac{f(x)}{g(x)}$ ，不妨令 $f(x),g(x)$ 都乘上 ${\prod }_{i=-S}^S(1+\frac{i}{S}x)$ ，这样 $f(x),g(x),f^{\prime }(x),g^{\prime }(x)$ 显然在 $x=1$ 处就都有实数取值了。
• 记新得到的 $$f(x)=\sum _{i=-S}^S{a}_i\prod _{j=-S,j̸=i}^S(1+\frac{j}{S}x)=a_{-S}\prod _{j=-S+1}^S(1+\frac{j}{S}x)$$
• 对 $f(x)$ 求导，得 $$f^{\prime }(x)=\sum _{i=-S}^S{a}_i\sum _{j=-S,j̸=i}^S\frac{j}{S}\prod _{k=-S,k̸=i,j}^S(1+\frac{k}{S}x)$$
• 计算上式时只需计算 $i=-S$ 或 $j=-S$ 的情况即可，其余情况 $\prod$ 内的值为 $0$ 。
• 时间复杂度 $O(NV+VLogV)$ 。

【代码】

#include<bits/stdc++.h>
using namespace std;
const int MAXN = 105;
const int MAXV = 1e5 + 5;
const int P = 998244353;
typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;
template <typename T> void chkmax(T &x, T y) {x = max(x, y); }
template <typename T> void chkmin(T &x, T y) {x = min(x, y); } 
template <typename T> void read(T &x) {
	x = 0; int f = 1;
	char c = getchar();
	for (; !isdigit(c); c = getchar()) if (c == '-') f = -f;
	for (; isdigit(c); c = getchar()) x = x * 10 + c - '0';
	x *= f;
}
template <typename T> void write(T x) {
	if (x < 0) x = -x, putchar('-');
	if (x > 9) write(x / 10);
	putchar(x % 10 + '0');
}
template <typename T> void writeln(T x) {
	write(x);
	puts("");
}
int n, sum, invs, s[MAXN], p[MAXN];
int f[MAXV], g[MAXV];
void update(int &x, int y) {
	x += y;
	if (x >= P) x -= P;
}
int power(int x, int y) {
	if (y == 0) return 1;
	int tmp = power(x, y / 2);
	if (y % 2 == 0) return 1ll * tmp * tmp % P;
	else return 1ll * tmp * tmp % P * x % P;
}
int func(int *f) {
	int ans = 1;
	for (int i = -sum + 1; i <= sum; i++)
		ans = 1ll * ans * (1 + (P + i + 0ll) * invs % P) % P;
	return 1ll * ans * f[0] % P;
}
int der(int *f) {
	static int tmp[MAXV];
	int ans = 0, tsum = 1;
	for (int i = -sum + 1; i <= sum; i++)
		tsum = 1ll * tsum * (1 + (P + i + 0ll) * invs % P) % P;
	for (int i = -sum + 1; i <= sum; i++)
		tmp[i + sum] = 1ll * tsum * power(1 + (P + i + 0ll) * invs % P, P - 2) % P;
	for (int i = -sum + 1; i <= sum; i++) {
		update(ans, 1ll * f[0] * (P + i) % P * invs % P * tmp[i + sum] % P);
		update(ans, 1ll * f[i + sum] * (P - sum) % P * invs % P * tmp[i + sum] % P);
	}
	return ans;
}
int main() {
	read(n);
	for (int i = 1; i <= n; i++)
		read(s[i]);
	for (int i = 1; i <= n; i++)
		read(p[i]), sum += p[i];
	int cur = 0; invs = power(sum, P - 2);
	f[sum] = g[sum] = 1;
	for (int i = 1; i <= n; i++) {
		static int tf[MAXV], tg[MAXV];
		memset(tf, 0, sizeof(tf));
		memset(tg, 0, sizeof(tg));
		for (int j = -cur; j <= cur; j++) {
			if (s[i] == 0) update(tf[j + sum - p[i]], f[j + sum]);
			else update(tf[j + sum - p[i]], P - f[j + sum]);
			update(tf[j + sum + p[i]], f[j + sum]);
			update(tg[j + sum - p[i]], g[j + sum]);
			update(tg[j + sum + p[i]], g[j + sum]);
		}
		cur += p[i];
		memcpy(f, tf, sizeof(tf));
		memcpy(g, tg, sizeof(tg));
	}
	reverse(f, f + 2 * sum + 1);
	reverse(g, g + 2 * sum + 1);
	int num = 0, den = 0;
	update(num, 1ll * func(g) * der(f) % P);
	update(num, P - 1ll * func(f) * der(g) % P);
	update(den, 1ll * func(g) * func(g) % P);
	writeln(1ll * num * power(den, P - 2) % P); 
	return 0;
}