期望题

题目链接：YBT2023寒假Day8 B

题目大意

给你一个 n 个节点的数，每个点有黑色或者白色。
你一开始在 1 号点，一直进行下面的操作：
如果点第一次到或者是黑色就把计数器加一，然后如果当前点度数为 1 就停止操作，否则等概率选择一个直接相连的点走过去。
问你结束之后计数器值的期望。
保证 1 号点度数大于 1。

思路

（首先停下来的条件是走到叶子没问题）
考虑黑色点白色点的贡献分开来算。

对于黑色：（那计数器就是步数，答案是期望步数）
考虑设 \(f_x\) 为 \(x\) 走到叶子的期望步数。
对于叶子有 \(f_x=[col_x=black]\)，那其他点就是 \(f_x=[col_x=black]+\dfrac{1}{d_x}\sum\limits_{(x,y)}f_y\)。

不过这个包含了父亲，自己，儿子三种，似乎不好 DP。
于是有一种方法是用主元思想，你考虑把 \(f_x\) 表示成 \(k_xf_{fa_x}+b_x\) 的形式。
那在最后 \(f_1\) 就是 \(b_1\) 了。

那带入进去：
\(f_x=[col_x=black]+\dfrac{1}{d_x}(f_{fa_x}+\sum\limits_{y\in son_x}(k_yf_x+b_y))\)
那你可以把 \(f_x\) 的所有项放到左边除就能得到与 \(f_{fa_x}\) 有关的式子了。

对于白色：
那你考虑每个点到它的概率，然后所有加起来就是白点的期望贡献。
那怎么到呢，一定是从父亲到。（第一次肯定要从父亲下来）
那设到的概率是 \(p_x\)，那就是 \(p_x=p_{fa_x}g_x\)。
（那这里是设 \(g_x\) 为从 \(fa_x\) 走到 \(x\) 的概率）

然后再推 \(g_x\)：
\(g_x=\dfrac{1}{d_{fa_x}}(1+g_{fa_x}g_x+\sum\limits_{y\in son_{fa_x},y\neq x}h_yg_x)\)
（三个分别是直接走到，走到父亲，走到别的儿子）

然后 \(h_x\) 就是设从 \(x\) 走到 \(fa_x\) 的概率。
那 \(h_x=\dfrac{1}{d_x}(1+\sum\limits_{y\in son_x}h_yh_x)\)

那这些式子都可以通过移项来解，然后就可以在树上 DP了。
（注意每个是从上往下还是从下往上要看清楚）

代码

#include<cstdio>
#include<vector>
#define mo 998244353

using namespace std;

const int N = 1e5 + 100;
int n, c[N], ans, du[N], fa[N];
vector <int> G[N];

int add(int x, int y) {return x + y >= mo ? x + y - mo : x + y;}
int dec(int x, int y) {return x < y ? x - y + mo : x - y;}
int mul(int x, int y) {return 1ll * x * y % mo;}
int ksm(int x, int y) {
	int re = 1;
	while (y) {
		if (y & 1) re = mul(re, x);
		x = mul(x, x); y >>= 1;
	}
	return re;
}

int k[N], b[N];

void dfs0(int now, int father) {
	if (du[now] == 1) {
		b[now] = (c[now] == 1); return ;
	}
	b[now] = (c[now] == 1) * du[now]; int num = du[now];
	k[now] = 1; fa[now] = father; 
	for (int i = 0; i < G[now].size(); i++) {
		int x = G[now][i]; if (x == father) continue;
		dfs0(x, now);
		num = dec(num, k[x]);
		b[now] = add(b[now], b[x]);
	}
	num = ksm(num, mo - 2);
	k[now] = mul(k[now], num); b[now] = mul(b[now], num);
}

void slove_black() {
	dfs0(1, 0);
	ans = add(ans, b[1]);
}

int p[N], g[N], h[N];

void dfs1(int now, int father) {
	if (G[now].size() == 1) return ;
	int num = du[now]; h[now] = 1;
	for (int i = 0; i < G[now].size(); i++) {
		int x = G[now][i]; if (x == father) continue;
		dfs1(x, now); num = dec(num, h[x]);
	}
	num = ksm(num, mo - 2);
	h[now] = mul(h[now], num);
}

void dfs2(int now, int father) {
	if (now != 1) {
		int num = du[father];
		g[now] = 1;
		num = dec(num, g[father]);
		for (int i = 0; i < G[father].size(); i++) {
			int x = G[father][i]; if (x == fa[father] || x == now) continue;
			num = dec(num, h[x]);
		}
		num = ksm(num, mo - 2);
		g[now] = mul(g[now], num);
	}
	if (now == 1) p[now] = 1;
		else p[now] = mul(p[father], g[now]);
	for (int i = 0; i < G[now].size(); i++) {
		int x = G[now][i]; if (x == father) continue;
		dfs2(x, now);
	}
}

void slove_white() {
	dfs1(1, 0);
	dfs2(1, 0);
	for (int i = 1; i <= n; i++) if (c[i] == 0) ans = add(ans, p[i]);
}

int main() {
	freopen("tree.in", "r", stdin);
	freopen("tree.out", "w", stdout);
	
	scanf("%d", &n);
	for (int i = 1; i <= n; i++) scanf("%1d", &c[i]);
	for (int i = 1; i < n; i++) {
		int x, y; scanf("%d %d", &x, &y);
		G[x].push_back(y); G[y].push_back(x);
		du[x]++; du[y]++;
	}
	
	slove_black();
	slove_white();
	printf("%d", ans);
	
	return 0;
}