CF1626E Black and White Tree

CF1626E Black and White Tree

换根 dp

树上路径行走问题，因其节点的转移不止于其子树有关，一般考虑换根 dp 或寻找新的转移顺序。

在这题里，考虑一个以 \(i\) 为点的子树，判断 \(i\) 是否可以走到子树中某个黑点，设 \(f_u\) 表示 \(u\) 能否走到黑点，枚举儿子 \(v\)，有三种满足方式：

1. \(a_u=1\)

2. \(a_v=1\)

3. \(f_{v}=1\) 且 \(sum_v>1\)

对于第三种方式，此时 \(u\) 已经求出了以 \(1\) 为根，进入其子树的答案，所以我们只需要考虑是否可以进入 \(v\) 子树中的黑点，那么需要满足 \(v\) 子树中至少有两个黑点才可以。可以套了换根 dp 更好理解。

那么这样的 dp 形式直接用换根 dp 即可。先求出以 \(1\) 为根，所有的 \(f_i\)，此时还没考虑全。然后用 \(1\) 节点拓展即可。

#include <bits/stdc++.h>
#define pii std::pair<int, int>
#define fi first
#define se second
#define pb push_back

typedef long long i64;
const i64 iinf = 0x3f3f3f3f, linf = 0x3f3f3f3f3f3f3f3f;
const int N = 3e5 + 10;
int n;
std::vector<int> e[N];
int a[N], sz[N], f[N], g[N];
void dfs(int u, int fa) {
	sz[u] = a[u];
	if(a[u]) f[u] = 1;
	for(auto v : e[u]) {
		if(v == fa) continue;
		if(a[v]) f[u] = 1;
		dfs(v, u);
		if(f[v] && sz[v] > 1) f[u] = 1;
		sz[u] += sz[v];
	}
}
void dfs2(int u, int fa) {
	if(f[u] || a[u]) g[u] = 1;
	for(auto v : e[u]) {
		if(v == fa) continue;
		if(a[u]) g[v] = 1;
		if(g[u] && sz[1] - sz[v] > 1) g[v] = 1;
		dfs2(v, u);
	}
}
void Solve() {
	std::cin >> n;
	for(int i = 1; i <= n; i++) std::cin >> a[i];

	for(int i = 1; i < n; i++) {
		int u, v;
		std::cin >> u >> v;
		e[u].pb(v), e[v].pb(u);
	}
	dfs(1, 0);
	dfs2(1, 0);
	for(int i = 1; i <= n; i++) {
		std::cout << g[i] << " \n"[i == n];
	}
}
int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    
	Solve();

	return 0;
}