JZOJ 6870. 【2020.11.17提高组模拟】ckw的树（树上期望DP+解方程）

JZOJ 6870. 【2020.11.17提高组模拟】ckw的树

题目大意

• 求大小为$N$树上每个点随机游走到$M$个标记点中任意一个的期望时间，每次等概率地到达与自己距离$\le 2$的点（包括自己）。
• $M,N\le 10^5$.

题解

• 暴力做法可以根据题意列出$N$条方程，高斯消元解出每个点的期望，时间复杂度$O(N^3)$。
• 这样做每个点的方程中会有它的祖父、父亲、兄弟层（包括自己）、儿子层、孙子层和常数，即为$E_k=a{E}_{f{a}_{f{a}_k}}+b{E}_{f{a}_k}+c\sum E_{bro}+d\sum E_{so{n}_k}+e\sum E_{so{n}_{so{n}_k}}+f$的形式，
• 其中每个点的祖父和父亲是唯一的，而兄弟、儿子、孙子会有若干，复杂度比较高。
• 起初，$a=b=c=d=e=\frac{1}{d_k}$，$f=1$，其中$d_k$为到$k$的距离$\le 2$的点的个数。
• 目标是把每个点的方程右边都消剩$a,b,c,f$项，
• 叶子节点处本身就只有$a,b,c,f$项，无需考虑。假设此时已经完成了所有后代的点，考虑如何把自己的右边也消剩$a,b,c,f$项。
• 也就是要消去$d,e$，即把$E_{son}$和$E_{so{n}_{son}}$的方程代入到自己的方程，
• 先考虑儿子，因为$c$项比较复杂，可以先把$c$项消去，
• 把每个儿子的方程相加，刚好得到$\sum E_{bro}=a^{\prime }{E}_{f{a}_{f{a}_k}}+b^{\prime }{E}_{f{a}_k}+c^{\prime }\sum E_{bro}+f^{\prime }$，其中$a^{\prime },b^{\prime },c^{\prime },f^{\prime }$分别为每个儿子$a,b,c,f$的和，于是发现$c$项可以移到左边，用$a,b,f$项表示，则$c$项被消去。
• 于是儿子只剩$a,b,f$项了，都可以直接代入到自己的方程中，此时儿子解决完了。
• 接着考虑孙子，类似的，在儿子处已经消剩了$a,b,f$项，那么也可以直接代入到自己的方程中，其中会涉及到儿子的期望，则再把儿子重复代入一遍，乘上对应的权值即可。
• 注意标记点的期望需默认为$0$。
• 最后每个点的方程之和父亲和祖父有关，而根节点是没有父亲和祖父的，则从上往下计算求出每个点的期望即可。

代码

• 由于实际操作中不需要$d,e$两项，所以代码中的$d$即为题解中的$f$。

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define N 100010
#define md 998244353
#define ll long long
int last[N], nxt[N * 2], to[N * 2], len = 0;
int p[N], d[N], si[N];
ll ans[N];
struct {
	ll a, b, c, d;
}f[N], g[N];
void add(int x, int y) {
	to[++len] = y;
	nxt[len] = last[x];
	last[x] = len;
}
ll ksm(ll x, ll y) {
	if(!y) return 1;
	ll l = ksm(x, y / 2);
	if(y % 2) return l * l % md * x % md;
	return l * l % md;
}
void dfs(int k, int fa, int ff) {
	d[k] = si[fa] + (fa > 0) + (ff > 0);
	for(int i = last[k]; i; i = nxt[i]) if(to[i] != fa) si[k]++;
	ll A = 0, B = 0, C = 1, D = 0;
	for(int i = last[k]; i; i = nxt[i]) if(to[i] != fa) {
		int x = to[i];
		dfs(x, k, fa);
		d[k] += si[x] + 1;
		A += f[x].a, B += f[x].b, C += md - f[x].c, D += f[x].d;
	}
	A %= md, B %= md, C %= md, D %= md;
	
	ll t = ksm(d[k], md - 2), K = 1;
	f[k].a = f[k].b = f[k].c = t, f[k].d = 1;
	ll s = ksm(C, md - 2);
	if(C == 0) s = 0;
	A = A * s % md, B = B * s % md, D = D * s % md;
	for(int i = last[k]; i; i = nxt[i]) if(to[i] != fa) {
		int x = to[i];
		f[x].a = (f[x].a + f[x].c * A) % md;
		f[x].b = (f[x].b + f[x].c * B) % md;
		f[x].d = (f[x].d + f[x].c * D) % md;
		f[x].c = 0;
		
		f[k].b = (f[k].b + f[x].a * (g[x].b + 1) % md * t) % md;
		K = (K - f[x].b * (g[x].b + 1) % md * t % md + md) % md;
		f[k].d = (f[k].d + f[x].d * (g[x].b + 1) % md * t % md) % md;
		
		f[k].d = (f[k].d + g[x].d * t) % md;
		K = (K - g[x].a * t % md + md) % md;
		
		g[k].a = (g[k].a + f[x].a) % md;
		g[k].b = (g[k].b + f[x].b) % md;
		g[k].d = (g[k].d + f[x].d) % md;
	}
	t = ksm(K, md - 2);
	f[k].a = f[k].a * t % md, f[k].b = f[k].b * t % md, f[k].c = f[k].c * t % md, f[k].d = f[k].d * t % md;
	if(p[k]) {
		f[k].a = f[k].b = f[k].c = f[k].d = 0;
	}
}
void solve(int k, int fa, int ff) {
	ans[k] = (f[k].d + ans[fa] * f[k].b + ans[ff] * f[k].a) % md;
	for(int i = last[k]; i; i = nxt[i]) if(to[i] != fa) {
		solve(to[i], k, fa);
	}
}
int main() {
	int n, m, i, j, k, x, y;
	scanf("%d%d", &n, &m);
	for(i = 1; i < n; i++) {
		scanf("%d%d", &x, &y);
		add(x, y), add(y, x);
	}
	for(i = 1; i <= m; i++) scanf("%d", &x), p[x] = 1;
	si[0] = 1;
	dfs(1, 0, 0);
	ll t = ksm((1 - f[1].c + md) % md, md - 2);
	f[1].a = f[1].a * t % md, f[1].b = f[1].b * t % md, f[1].d = f[1].d * t % md, f[1].c = 0;
	solve(1, 0, 0);
	for(i = 1; i <= n; i++) {
		printf("%lld
", ans[i]);
	}
	return 0;
}