题解：P5405 [CTS2019] 氪金手游

题意

给定一棵 $n$ 个点的树，边是有向的；现在不断从 $\set{1,2,\cdots,n}$ 随机抽取一个数（如果抽到之前抽到过的就重新抽），每个数被抽到的概率按照下面的方法确定：

• 首先，每个数有一个权值 $w_i$，一开始 $w_i$ 取值为 $j$ 的概率是 $P_{i,j}$，其中 $j\in[1,3]$；
• 每一轮抽到 $i$ 的概率是 $p_i=\frac{w_i}{\sum_{j=1}^nw_j}$。

求对于所有边 $u\to v$，都满足 $u$ 先于 $v$ 抽到的概率。对 $998244353$ 取模。

$n\le 1000$。

解法

首先，边的方向不是统一的，但是如果我们假设所有边的方向都是由父亲指向儿子，那么限制就变成：对于每个点 $u$，要求其余所有在 $u$ 子树 $Tr_u$ 中的点 $v\in Tr_u,v\ne u$ 都晚于 $u$ 被抽到；那么概率就相当于是在 $Tr_u$ 这些点中第一发就出了 $u$，即 $$ \frac{p_u}{\sum_{v\in Tr_u}p_v} $$ 整个的概率就是 $$ \prod_{u=1}^n\frac{p_u}{\sum_{v\in Tr_u}p_v} $$ 现在，其中有一些边不满足由父亲指向儿子。对于一个边集 $E’$，设 $g(E’)$ 表示对于一条边 $u\to v$，$u$ 先于 $v$ 抽到当且仅当 $u\to v\in E’$ 概率；转部分满足，设 $f(E’)$ 表示要求 $E’$ 中的每条边 $u\to v$ 都满足 $u$ 先于 $v$ 抽到、其余边无所谓的概率。那么考虑现在整棵树就会被不在 $E’$ 中的边划分成若干个连通块，限制从一个根 $u$ 先于整个子树 $Tr_u$ 变成先于 $u$ 所在连通块中以 $u$ 为根的部分，设这个点集是 $S_u$，那么也能够得到： $$ f(E’)=\prod_{u=1}^n\frac{p_u}{\sum_{v\in S_u}p_v} $$ 设 $E$ 表示原树中父亲指向儿子的边的边集，根据容斥原理有 $$ g(E)=\sum_{E\subseteq T}(-1)^{|T|-|E|}\prod_{u=1}^n\frac{p_u}{\sum_{v\in S_u}p_v} $$ 现在考虑 $p_u$ 也是随机变量，但其实可以直接把 $p_u$ 看成 $w_u$。先把 $(-1)^{|E|}$ 这个常量提出来，暴力枚举所有 $w$ 取值： $$ \begin{aligned}\sum_{E\subseteq T}(-1)^{|T|}\sum_{w_i=1}^3\left(\prod_{i=1}^nP_{i,w_i}\right)\prod_{u=1}^n\frac{\frac{w_u}{\sum_{i=1}^nw_i}}{\sum_{v\in S_u}\frac{w_v}{\sum_{i=1}^nw_i}}&=\sum_{E\subseteq T}(-1)^{|T|}\sum_{w_i=1}^3\left(\prod_{i=1}^nP_{i,w_i}\right)\prod_{u=1}^n\frac{w_u}{\sum_{v\in S_u}w_v}\&=\sum_{E\subseteq T}(-1)^{|T|}\sum_{w_i=1}^3\prod_{u=1}^n\frac{P_{u,w_u}w_u}{\sum_{v\in S_u}w_v}\end{aligned} $$ 考虑用 dp 计算这个容斥的式子。由于是在一棵树上做，因此自然考虑对每个点为根的一个子树设计状态，计算点集为一个子树时的贡献。先分离根的贡献，对于一个根节点 $r$，可以考虑先枚举它自己 $w$ 的取值为 $j$，以及其连通块 $\sum_{v\in S_r}w_v$ 的值 $k$；那么就有 $$ \sum_{w_i=1}^3\sum_{E\subseteq T}(-1)^{|T|}\prod_{u\in Tr_r}\frac{P_{u,w_u}w_u}{\sum_{v\in S_u}w_v}=\frac{P_{r,j}j}{k}\sum_{i\ne r,w_i=1}^3\sum_{E\subseteq T}\left\sum_{v\in S_r}w_v=k\right^{|T|}\prod_{u\in Tr_r,u\ne r}\frac{P_{u,w_u}w_u}{\sum_{v\in S_u}w_v} $$ 令 $\mathrm{son}u$ 表示点 $u$ 的所有儿子集合，那么对于所有 $v\in \mathrm{son}u$，可以将枚举 $r$ 子树的所有点变成枚举每个 $v$ 的子树中的点，于是可以进一步得到原式为 $$ \frac{Pj}{k}\sum{i\ne r,w_i=1}^3\sum_{E\subseteq T}\left\sum_{v\in S_r}w_v=k\right^{|T|}\prod_{u\in\mathrm{son}r}\prod{v\in Tr_u}\frac{P_{v,w_v}w_v}{\sum_{x\in S_v}w_x} $$ 这样就将连乘的部分分离到了各个孩子的子树，现在处理对所有儿子的 $w_v$ 的和的限制。设在边集为 $T$ 的情况下，$u$ 的所有儿子中与 $u$ 在同一连通块的儿子集合为 $\mathrm{son}^u$，那么 $S_r$ 可以看成这些儿子的 $S$ 的并集再加上 $r$，即 $$ S_r=\set{r}\cup \left(\bigcup{u\in \mathrm{son}^r}S_u\right) $$ 那么上面 $\left(\sum{v\in S_r}w_v\right)=k$ 的限制就可以写成 $$ w_r+\sum_{u\in\mathrm{son}^r}\sum{v\in S_u}w_v=k\Longrightarrow\sum_{u\in\mathrm{son}^r}\sum{v\in S_u}w_v=k-j $$ 令 $k_u=\sum_{v\in S_u}w_v$，那么原式相当于要求所有与 $r$ 在同一连通块的儿子的 $k_u$ 之和为 $k-j$。考虑枚举 $r$ 的每个儿子 $u$ 的 $k_u$ 的值，对于在 $\mathrm{son}^r$ 中的儿子的 $k_u$ 的和有限制，其余儿子无限制；那么原式等于 $$ \frac{Pj}{k}\sum_{i\ne r,w_i=1}^3\sum_{E\subseteq T}(-1)^{|T|}\sum_{\sum_{u\in\mathrm{son}^r}k_u=k-j}\prod{u\in\mathrm{son}r}\left[\sum{v\in S_u}w_v=k_u\right]\prod_{v\in Tr_u}\frac{P_{v,w_v}w_v}{\sum_{x\in S_v}w_x} $$ 再考虑 $(-1)^{|T|}$ 如何分离。设 $T_u$ 表示所有在 $u$ 的子树中的、从父亲指向儿子的边的边集，那么根据 $r$ 和每个儿子是否在同一连通块，显然有 $$ T_r=\left(\bigcup_{u\in{\mathrm{son}^r}}r\to u\right)\cup\left(\bigcup{u\in\mathrm{son}}T_u\right) $$ 于是可以分离 $|T_r|=|\mathrm{son}^r|+\sum{u\in\mathrm{son}u}|T_u|$，那么原式等于 $$ \frac{Pj}{k}\sum_{i\ne r,w_i=1}^3\sum_{E\subseteq T}(-1)^{|\mathrm{son}^r|}\sum{\sum_{u\in\mathrm{son}^r}k_u=k-j}\prod{u\in\mathrm{son}r}(-1)^{|T_u|}\left[\sum{v\in S_u}w_v=k_u\right]\prod_{v\in Tr_u}\frac{P_{v,w_v}w_v}{\sum_{x\in S_v}w_x} $$ 再考虑 $E\subseteq T$ 实际上可以看成：先在 $r$ 与每个儿子的连边中选一个超集放到 $T$ 里，再在每个子树的部分选一个超集放到 $T$ 里。设 $E_u$ 表示 $E$ 中在 $u$ 的子树里的部分的边集，$\mathrm{Eson}^{}u$ 表示边集为 $E$ 的情况下，也就是原树中与 $u$ 在同一连通块的儿子的集合。顺便把最开始对子树中每个点 $w_i$ 的取值的枚举也放进各个孩子的子树里，那么原式等于 $$ \frac{Pj}{k}\sum_{\mathrm{Eson}^r\subseteq \mathrm{son}^_r}(-1)^{|\mathrm{son}^r|}\sum{\sum{u\in\mathrm{son}^r}k_u=k-j}\prod{u\in\mathrm{son}r}\sum{i\in Tr_u,w_i\in[1,3]}\sum_{E_u\subseteq T_u}(-1)^{|T_u|}\left[\sum_{v\in S_u}w_v=k_u\right]\prod_{v\in Tr_u}\frac{P_{v,w_v}w_v}{\sum_{x\in S_v}w_x} $$ 来看一眼一开始的式子： $$ \sum_{E\subseteq T}(-1)^{|T|}\sum_{w_i=1}^3\prod_{u=1}^n\frac{P_{u,w_u}w_u}{\sum_{v\in S_u}w_v} $$ 注意到从第 $3$ 个 $\sum$ 开始，后面的式子就是在计算整棵树为 $Tr_u$ 时的答案。现在考虑令 $F(u,k)$ 表示在 $k_u=k$ 的情况下子树 $u$ 的答案，那么有 $$ F(r,k)=\sum_{j=1}^3\frac{P_{r,j}j}{k}\sum_{\mathrm{Eson}^r\subseteq \mathrm{son}^_r}(-1)^{|\mathrm{son}^r|}\sum{\sum{u\in \mathrm{son}^r}k_u=k-j}\prod{u\in \mathrm{son}_r}F(u,k_u) $$ 现在考虑 $\mathrm{Eson}^r$ 的超集和枚举。对于一个孩子 $u\in\mathrm{son}r$，如果 $u\not\in \mathrm{son}^*r$ 就说明 $u$ 对不会受 $k-j$ 的限制，且这些儿子的贡献都是独立的。于是就有： $$ F(r,k)=\sum^3\frac{Pj}{k}\sum{\mathrm{Eson}^_r\subseteq \mathrm{son}^r}\left(\sum{\sum_{u\in \mathrm{son}^r}k_u=k-j}\prod{u\in \mathrm{son}^r}-F(u,k_u)\right)\left(\prod{u\in\mathrm{son}r,u\not\in\mathrm{son}^*r}\sumF(u,k_u)\right) $$ 限制若干项总和是多少，然后将这些项乘起来，可以考虑使用生成函数：令 $f_r=\sum{k\ge 0}F(r,k)x^k$，那么就有 $$ [x^k]f_r=\sum_{j=1}^3\frac{P_{r,j}j}{k}[x^{k-j}]\sum_{\mathrm{Eson}r^\subseteq\mathrm{son}_r^}\left(\prod{u\in\mathrm{son}r^*}-f_u\right)\left(\prod{u\in\mathrm{son}r,u\not\in\mathrm{son}^_r}f_u(1)\right) $$ 最后，对于 $u\in\mathrm{son}_r$，如果 $u\in\mathrm{Eson}^r$ 那么说明 $u$ 对答案只能有一种贡献：$-f_u$；否则会有两种贡献：$-f_u,f_u(1)$；使用乘法分配律，即可得到最终答案的表达式： $$ [x^k]f_r=\sum^3\frac{Pj}{k}[x^{k-j}]\left(\prod_{u\in\mathrm{Eson}r^*}-f_u\right)\left(\prod{u\in\mathrm{son}r,u\not\in\mathrm{Eson}^*r}f_u(1)-f_u\right) $$ 当 $r$ 是叶子节点时，显然 $f_r=Px+Px^2+P_{r,3}x^3$。注意到 $f_r$ 的次数界是 $O(\text{子树大小})$ 的，使用平方暴力计算卷积，根据树上背包的结论时间复杂度 $O(n^2)$，可以通过本题。不要忘了一开始提出了一个常数项 $(-1)^{|E|}$。

代码

#include <bits/stdc++.h>
bool MemoryST; using namespace std;
#define ll long long
#define mk make_pair
#define open(x) freopen(#x".in", "r", stdin), freopen(#x".out", "w", stdout)
#define lowbit(x) ((x) & (-(x)))
#define lson l, mid, rt << 1
#define rson mid + 1, r, rt << 1 | 1
#define BCNT __builtin_popcount
#define cost_time (1e3 * clock() / CLOCKS_PER_SEC) << "ms"
#define cost_space (abs(&MemoryST - &MemoryED) / 1024.0 / 1024.0) << "MB"
const int inf = 0x3f3f3f3f; 
const ll linf = 1e18; 
mt19937 rnd(random_device{}());
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> T abs(T x) { return (x < 0) ? -x : x; }
const int P = 998244353;
namespace ModOptions_normal {
    int add(int x, int y) { return (x + y) % P; } void addto(int &x, int y) { (x += y) %= P; }
    int mul(int x, int y) { return 1ll * x * y % P; } void multo(int &x, int y) { x = mul(x, y); }
    int qp(int x, int y) {
        int res = 1; for (; y; y >>= 1, multo(x, x))
            if (y & 1) multo(res, x);
        return res;
    } int divv(int x, int y) { return 1ll * x * qp(y, P - 2) % P; } void divto(int &x, int y) { x = divv(x, y); } 
    int sub(int x, int y) { return (x + P - y) % P; } void subto(int &x, int y) { x = sub(x, y); }
} using namespace ModOptions_normal;
bool MemoryED; int main() {
    int n; cin >> n; vector<array<int, 3> > P(n);
    for (int i = 0; i < n; i ++) {
        int a, b, c; cin >> a >> b >> c; int inv = divv(1, add(add(a, b), c));
        P[i][0] = mul(a, inv), P[i][1] = mul(b, inv), P[i][2] = mul(c, inv);
    } vector<vector<pair<int, int> > > mp(n);
    int cnt = 0;
    for (int i = 1, u, v; i < n; i ++) {
        cin >> u >> v, u --, v --;
        mp[u].emplace_back(v, 0), mp[v].emplace_back(u, 1);
    } vector<vector<int> > f(n);
    auto poly_mul = [&](vector<int> &a, vector<int> &b) -> vector<int> {
        vector<int> c(a.size() + b.size() - 1, 0); int len1 = (int)a.size(), len2 = (int)b.size();
        for (int i = 0; i < len1; i ++)
            for (int j = 0; j < len2; j ++)
                addto(c[i + j], mul(a[i], b[j]));
        return c;
    }; auto getsum = [&](vector<int> &a) -> int {
        int res = 0, len = (int)a.size();
        for (int i = 0; i < len; i ++) addto(res, a[i]);
        return res;
    }; vector<int> fuyi; fuyi.resize(1, :: P - 1);
    auto dfs = [&](auto self, int u, int fa) -> void {
        vector<int> tmp(1, 1); bool is_leaf = 1;
        for (auto [v, tp] : mp[u]) if (v != fa) {
            is_leaf = 0, self(self, v, u); int sum = getsum(f[v]);
            f[v] = poly_mul(f[v], fuyi);
            if (tp == 1) f[v][0] = add(sum, f[v][0]);
            else cnt ++;
            tmp = poly_mul(tmp, f[v]);
        } if (is_leaf) return f[u].resize(4), f[u][1] = P[u][0], f[u][2] = P[u][1], f[u][3] = P[u][2], void(0);
        int siz = (int)tmp.size(); f[u].resize(siz + 3, 0);
        for (int j = 1; j <= 3; j ++)
            for (int k = j; k < j + siz; k ++)
                addto(f[u][k], mul(divv(mul(j, P[u][j - 1]), k), tmp[k - j]));
    }; dfs(dfs, 0, -1);
    printf("%d\n", (cnt & 1) ? sub(0, getsum(f[0])) : getsum(f[0]));
    return 0;
}