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;
}