题目大意
给定 n n n 个节点的无根树,有 k k k 个指定的节点上有蝴蝶。对每条边等概率地指定任意方向,按照边的编号顺序查找每条边,若边 u → v u\to v u→v 满足 u u u 有蝴蝶且 v v v 没有蝴蝶,则蝴蝶从 u u u 飞到 v v v. 求两只蝴蝶之间距离的期望值。
题解
首先要注意,求期望值的两只蝴蝶不是给定的。
显然,这个期望值的求法是每条边的贡献累加,求出所有蝴蝶两两距离和的期望值,再除以选择两只蝴蝶的方案数 C k 2 = k ( k − 1 ) 2 C_k^2=\frac {k(k-1)}{2} Ck2=2k(k−1).
先考虑一个简化的经典问题:蝴蝶不会飞的情况。
枚举每条边,若其左边有 x x x 只蝴蝶,那么其右边有 k − x k-x k−x 只蝴蝶,这条边将会成为 x × ( k − x ) x\times (k-x) x×(k−x) 对蝴蝶之间的路径上的一条边。换一种说法,有 x × ( k − x ) x\times (k-x) x×(k−x) 条路径会经过这条边。那么这条边会对总距离值贡献 x × ( k − x ) x\times (k-x) x×(k−x). 也就是说简化情况下 a n s = ∑ ( u , v ) ∈ M s u m u ( k − s u m u ) k ( k − 1 ) 2 ans=\frac{\sum\limits_{(u,v)\in M}sum_u(k-sum_u)}{\frac{k(k-1)}{2}} ans=2k(k−1)(u,v)∈M∑sumu(k−sumu). s u m i sum_i sumi 表示 i i i 的子树中有多少只蝴蝶。这个可以一次dfs O ( n ) O(n) O(n) 求出。
回到原问题。如果我们知道每条边两边分别有多少蝴蝶,则贡献可求。
思考发现对于边 ( u , v ) (u,v) (u,v),不管蝴蝶怎么飞, s u m sum sum 值的变化量最大为 1 1 1. 而 s u m sum sum 是否改变只取决于 u u u 和 v v v 上是否有蝴蝶。这就是说,我们需要求出 u u u 和 v v v 有蝴蝶的概率。设这个概率为 p p p.
考虑对于 p u p_u pu:
- 原来 u u u 和 v v v 上都有蝴蝶,则现在 u u u 有蝴蝶的概率为 p u × p v p_u\times p_v pu×pv.
- 原来 u u u 和 v v v 上都没有蝴蝶,则现在 u u u 有蝴蝶的概率为 0 0 0.
- 原来 u u u 有蝴蝶, v v v 没有,则现在 u u u 有蝴蝶的概率为 p u ( 1 − p v ) 2 \frac {p_u(1-p_v)}{2} 2pu(1−pv). 因为有 1 2 \frac 1 2 21 的概率边 ( u , v ) (u,v) (u,v) 变成 v → u v\to u v→u,只有这样蝴蝶才不会从 u u u 飞到 v v v.
- 原来 u u u 没有蝴蝶, v v v 有,则现在 u u u 有蝴蝶的概率为 ( 1 − p u ) p v 2 \frac {(1-p_u)p_v}{2} 2(1−pu)pv. 因为有 1 2 \frac 1 2 21 的概率边 ( u , v ) (u,v) (u,v) 变成 v → u v\to u v→u,只有这样蝴蝶才会从 v v v 飞到 u u u.
对上述概率求和,得到 u u u 有蝴蝶的概率为 p u p v + p u ( 1 − p v ) 2 + ( 1 − p u ) p v 2 = p u + p v 2 p_up_v+\frac {p_u(1-p_v)}{2}+\frac {(1-p_u)p_v}{2}=\frac {p_u+p_v}{2} pupv+2pu(1−pv)+2(1−pu)pv=2pu+pv.
也就是转移时 p u = p u + p v 2 p_u=\frac {p_u+p_v}{2} pu=2pu+pv. 对于 p v p_v pv 同理,得 p u = p v = p u + p v 2 p_u=p_v=\frac {p_u+p_v}{2} pu=pv=2pu+pv.
在修改概率前,我们就需要对边 ( u , v ) (u,v) (u,v) 统计答案。
考虑如下情况:
- u u u 和 v v v 都有蝴蝶,贡献为 p u ⋅ p v ⋅ s u m u ⋅ ( k − s u m u ) p_u\cdot p_v\cdot sum_u\cdot (k-sum_u) pu⋅pv⋅sumu⋅(k−sumu)
- u u u 和 v v v 都没蝴蝶,贡献为 ( 1 − p u ) ⋅ ( 1 − p v ) ⋅ s u m u ⋅ ( k − s u m u ) (1-p_u)\cdot (1-p_v)\cdot sum_u\cdot (k-sum_u) (1−pu)⋅(1−pv)⋅sumu⋅(k−sumu)
- u u u 有蝴蝶, v v v 没有蝴蝶,此时有 1 2 \frac 1 2 21 的概率蝴蝶从 u u u 飞到 v v v,导致 s u m u sum_u sumu 少 1 1 1,贡献为 p u ⋅ ( 1 − p v ) ⋅ s u m u ⋅ ( k − s u m u ) + ( s u m u − 1 ) ( k − s u m u + 1 ) 2 p_u\cdot (1-p_v)\cdot \frac {sum_u\cdot (k-sum_u)+(sum_u-1)(k-sum_u+1)}{2} pu⋅(1−pv)⋅2sumu⋅(k−sumu)+(sumu−1)(k−sumu+1)
- u u u 没有蝴蝶, v v v 有蝴蝶,此时有 1 2 \frac 1 2 21 的概率蝴蝶从 v v v 飞到 u u u,导致 s u m u sum_u sumu 多 1 1 1,贡献为 ( 1 − p u ) ⋅ p v ⋅ s u m u ⋅ ( k − s u m u ) + ( s u m u + 1 ) ( k − s u m u − 1 ) 2 (1-p_u)\cdot p_v\cdot \frac {sum_u\cdot (k-sum_u)+(sum_u+1)(k-sum_u-1)}{2} (1−pu)⋅pv⋅2sumu⋅(k−sumu)+(sumu+1)(k−sumu−1)
对以上 4 4 4 种情况求和就是边 ( u , v ) (u,v) (u,v) 的贡献。
具体实现的时候可以对 u u u 和 v v v 限定父子关系,这样容易计算。
记得最后要除以 k ( k − 1 ) 2 \frac {k(k-1)}{2} 2k(k−1).
代码
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 300005;
const LL mod = 998244353;
int n, k, cnt = 0, fir[N], nxt[N << 1], to[N << 1], fa[N];
LL ans = 0, p[N], s[N];
void ade(int u, int v) {
cnt++, nxt[cnt] = fir[u], fir[u] = cnt, to[cnt] = v;
cnt++, nxt[cnt] = fir[v], fir[v] = cnt, to[cnt] = u;
}
LL qpow(LL x, LL y) {
LL sum = 1;
while (y) {
if (y & 1) sum = sum * x % mod;
x = x * x % mod, y >>= 1;
}
return sum;
}
void dfs(int r, int f) {
fa[r] = f, s[r] = p[r];
for (int i = fir[r]; i; i = nxt[i])
if (to[i] != f)
dfs(to[i], r), s[r] += s[to[i]];
}
int main() {
scanf("%d%d", &n, &k);
for (int i = 1, x; i <= k; i++) scanf("%d", &x), p[x] = 1;
for (int i = 1, u, v; i < n; i++) scanf("%d%d", &u, &v), ade(u, v);
dfs(1, 0);
LL inv2 = qpow(2, mod - 2);
for (int i = 1; i <= cnt; i += 2) {
int u = to[i], v = to[i + 1];
if (fa[v] != u) swap(u, v);//限定 u 为 v 的父亲
ans = ((ans
+ p[u] * p[v] % mod * s[v] % mod * (k - s[v]) % mod
+ (1 - p[u]) * (1 - p[v]) % mod * s[v] % mod * (k - s[v]) % mod
+ (1 - p[u]) * p[v] % mod * (s[v] * (k - s[v]) % mod + (s[v] - 1) * (k - s[v] + 1) % mod) % mod * inv2 % mod
+ p[u] * (1 - p[v]) % mod * (s[v] * (k - s[v]) % mod + (s[v] + 1) * (k - s[v] - 1) % mod) % mod * inv2 % mod) % mod + mod) % mod;
p[u] = p[v] = (p[u] + p[v]) * inv2 % mod;//概率转移
}
printf("%lld", (ans * qpow(1ll * k * (k - 1) / 2 % mod, mod - 2) % mod + mod) % mod);
return 0;
}