LOJ 3045: 洛谷 P5326: 「ZJOI2019」开关

题目传送门：LOJ #3045。

题意简述

略。

题解

从高斯消元出发好像需要一些集合幂级数的知识，就不从这个角度思考了。

---

令 \(\displaystyle \dot p = \sum_{i = 1}^{n} p_i\)。

我们考虑一个操作序列 \(\{a_1, a_2, \ldots , a_k\}\)，其中 \(1 \le a_j \le n\)，就表示第 \(i\) 次按下了开关 \(a_j\)。

那么按 \(k\) 次后恰好得到这个序列的概率就是 \(\displaystyle \prod_{j = 1}^{k} (p_{a_j} / \dot p)\)。

那么我们考虑如果按下这个序列后恰好得到了目标状态 \(s\)：

当且仅当对于每个 \(i\)（\(1 \le i \le n\)）均满足按下开关 \(i\) 的次数的奇偶性恰好等于 \(s_i\)。

形式化地说，就是对于每个 \(i\) 有 \(\displaystyle \left( \sum_{j = 1}^{k} [a_j = i] \right) \bmod 2 = s_i\)。

---

那么我们对每个 \(i\) 分开考虑，对于 \(s_i = 0\) 的需要按偶数次，对于 \(s_i = 1\) 的需要按奇数次。

• 对于某个 \(s_i = 0\) 的 \(i\)，我们给出这样的数列：\(f_i = \{1, 0, {(p_i / \dot p)}^2, 0, {(p_i / \dot p)}^4, 0, {(p_i / \dot p)}^6, 0, \ldots \}\)。
• 对于某个 \(s_i = 1\) 的 \(i\)，我们给出这样的数列：\(f_i = \{0, p_i / \dot p, 0, {(p_i / \dot p)}^3, 0, {(p_i / \dot p)}^5, 0, {(p_i / \dot p)}^7, \ldots \}\)。

可以发现，把所有的 \(i\) 的数列全部二项卷积起来，就得到了一个新的数列 \(f\)，这个数列满足：

对于 \(f\) 的第 \(k\) 项 \(f_k\)，就表示了当按 \(k\) 下开关时，恰好得到状态 \(s\) 的概率。

因为是 二项卷积，所以我们把这个过程写成 指数型概率生成函数 的形式：

定义 \(\hat F_i (x) = \mathbf{EGF} \left( { \left\{ [j \bmod 2 = s_i] {(p_i / \dot p)}^j \right\} }_{j = 0}^{\infty} \right)\)，

也就是每个 \(i\) 对应的上述数列 \(f_i\) 的指数型生成函数，

写做封闭形式，就是 \(\displaystyle \hat F_i (x) = \frac{e^{(p_i / \dot p) x} + {(-1)}^{s_i} e^{-(p_i / \dot p) x}}{2}\)。

所以最终得到的 \(f\) 的 EGF 就是 \(\displaystyle \hat F (x) = \prod_{i = 1}^{n} \frac{e^{(p_i / \dot p) x} + {(-1)}^{s_i} e^{-(p_i / \dot p) x}}{2}\)。

看起来非常的变态，但是还没完！出什么问题了？

---

首先我们要明确：得到 \(f\) 能干啥？

发现 \(f\) 的性质是：\(f_k\) 表示按恰好 \(k\) 次开关得到状态 \(s\) 的概率，那么根据期望的定义，答案就是 \(\displaystyle \sum_{i = 0}^{\infty} i f_i\)。

这是啥啊，就是 \(f\) 对应的 普通生成函数 \(\displaystyle F(x) = \sum_{i = 0}^{\infty} f_i x^i\)，它在 \(1\) 处的导数，也就是 \(F'(1)\)。

（回顾形式幂级数求导，以及求值的定义）

但是 错了，再观察一下，题目要求的是 第一次 到达状态 \(s\) 的期望步数，而不是现在这个样子。

（因为可能不是第一次，而是此前已经经过很多次了。实际上如果直接求这个，甚至是不收敛的）

那么怎么办呢？我们发现需要排除第一次到达 \(s\) 后，又经过若干步返回 \(s\) 的情况，也就是返回原状态了。

由此，我们考虑求出数列 \(g\)，其中 \(g_k\) 表示在 \(k\) 步后恰好返回原状态的概率。

那么可以发现，如果令最终答案的数列为 \(h\)，有 \(h \ast g = f\)（\(h\) 卷 \(g\) 等于 \(f\)，是普通卷积不是二项卷积）。

而 \(g\) 应该如何求得呢？其实就是当全部 \(s_i = 0\) 时的 \(f\) 啦，因为是要返回原状态嘛。

---

上面说了一堆理论上的东西，现在我们考虑如何实现。

首先发现求的时候是 EGF，但是算答案的时候是 OGF，这很怪。我们观察一下形式看看能不能转换。

对于 \(\hat F\)，有形式 \(\displaystyle \hat F (x) = \prod_{i = 1}^{n} \frac{e^{(p_i / \dot p) x} + {(-1)}^{s_i} e^{-(p_i / \dot p) x}}{2}\)。

我们把每个形如 \(a_w e^{(w / \dot p)x}\) 的式子看作一项，可以发现最终 \(w\) 的取值在 \([-\dot p, \dot p]\)。

所以把 \(\hat F (x)\) 表示成 \(\displaystyle \sum_{w = -\dot p}^{\dot p} a_w e^{(w / \dot p) x}\) 的形式后，我们就有 \(\displaystyle f_k = \sum_{w = -\dot p}^{\dot p} a_w {(w / \dot p)}^k\)。

再把这个形式转换成 OGF，得到 \(\displaystyle \mathbf{OGF} (f) = F(x) = \sum_{w = -\dot p}^{\dot p} \frac{a_w}{1 - (w / \dot p) x}\)。

这时候考虑求出每个 \(a_w\)，可以发现做一个背包就行了，复杂度为 \(\mathcal O (n \dot p)\)。

（观察背包转移时的系数都是 \(\pm 1 / 2\)，可以使用多项式 Exp 优化到 \(\mathcal O (n + \dot p \log \dot p)\)，但是没有必要）

对于 \(\displaystyle \mathbf{OGF} (g) = G(x) = \sum_{w = -\dot p}^{\dot p} \frac{b_w}{1 - (w / \dot p) x}\) 同理，我们需要求出每一个 \(b_w\)。

求出所有 \(a_w, b_w\) 之后，我们就掌握了 \(f, g\) 的一些性质，然后对于答案 \(h\)，令其普通生成函数为 \(H\)。

则根据上面的解释，有 \(H = F / G\)，并且最终我们需要求出 \(H'(1)\)。

因为这里 \(F, G, H\) 都可能有无限项，所以要考虑通过 \(a_w, b_w\) 去求出答案。

考虑除法求导法则：\(\displaystyle H' = {(F / G)}' = \frac{F'G - G'F}{G^2}\)。

所以只要求出 \(F(1), G(1), F'(1), G'(1)\) 即可。

然而很可惜，我们发现因为存在 \(\displaystyle \frac{a_{\dot p}}{1 - (\dot p / \dot p) x}\) 这一项，所以 \(F, G, F', G'\) 在 \(x = 1\) 处不收敛。

我们知道答案一定收敛，所以考虑洛都可以洛做一点变换：把 \(F\) 和 \(G\) 都乘上 \((1 - x)\)。那么就有：

• \(F(1) = a_{\dot p}\)。
• \(\displaystyle F'(1) = \sum_{w = -\dot p}^{\dot p - 1} \frac{a_w}{w / \dot p - 1}\)。
• \(G(1) = b_{\dot p}\)。
• \(\displaystyle G'(1) = \sum_{w = -\dot p}^{\dot p - 1} \frac{b_w}{w / \dot p - 1}\)。

具体计算过程省略，就是按照求导的公式算而已。所以：

\[\begin{aligned} H'(1) &= \frac{F'(1) G(1) - G'(1) F(1)}{G^2(1)} \\ &= \frac{\displaystyle \left( \sum_{w = -\dot p}^{\dot p - 1} \frac{a_w}{w / \dot p - 1} \right) b_{\dot p} - \left( \sum_{w = -\dot p}^{\dot p - 1} \frac{b_w}{w / \dot p - 1} \right) a_{\dot p}}{b_{\dot p}^2} \\ &= \sum_{w = -\dot p}^{\dot p - 1} \frac{a_w b_{\dot p} - b_w a_{\dot p}}{(w / \dot p - 1) b_{\dot p}^2} \end{aligned}
\]

求出所有的 \(a_w, b_w\) 后按照此式计算即可，时间复杂度为 \(\mathcal O (n \dot p + \dot p \log mod)\)，代码如下：

#include <cstdio>
#include <algorithm>

typedef long long LL;
const int Mod = 998244353, Inv2 = (Mod + 1) / 2;
const int MN = 105, MP = 50005;

inline void Add(int &x, LL y) { x = (x + y) % Mod; }
inline int qPow(int b, int e) {
	int a = 1;
	for (; e; e >>= 1, b = (LL)b * b % Mod)
		if (e & 1) a = (LL)a * b % Mod;
	return a;
}
inline int gInv(int b) { return qPow(b, Mod - 2); }

int N, s[MN], p[MN], sump;
int _a[2][MP * 2], _b[2][MP * 2], *a[2] = {_a[0] + MP, _a[1] + MP}, *b[2] = {_b[0] + MP, _b[1] + MP};
int Ans;

int main() {
	scanf("%d", &N);
	for (int i = 1; i <= N; ++i) scanf("%d", &s[i]), s[i] = s[i] ? -1 : 1;
	b[0][0] = a[0][0] = 1;
	for (int i = 1; i <= N; ++i) {
		scanf("%d", &p[i]);
		for (int j = -sump - p[i]; j <= sump + p[i]; ++j) b[1][j] = a[1][j] = 0;
		for (int j = -sump; j <= sump; ++j)
			Add(a[1][j + p[i]], (LL)Inv2 * a[0][j]),
			Add(a[1][j - p[i]], s[i] * (LL)Inv2 * a[0][j]),
			Add(b[1][j + p[i]], (LL)Inv2 * b[0][j]),
			Add(b[1][j - p[i]], (LL)Inv2 * b[0][j]);
		sump += p[i];
		std::swap(a[0], a[1]), std::swap(b[0], b[1]);
	}
	int isump = gInv(sump), *A = a[0], *B = b[0];
	for (int j = -sump; j < sump; ++j)
		Add(Ans, ((LL)A[j] * B[sump] - (LL)B[j] * A[sump]) % Mod * gInv((LL)j * isump % Mod - 1));
	Ans = (LL)Ans * qPow(B[sump], Mod - 3) % Mod;
	printf("%d\n", (Ans + Mod) % Mod);
	return 0;
}