CF1442E Black, White and Grey Tree

CF1442E Black, White and Grey Tree

所有连在一起，颜色一样的点缩成一个点，这样相邻异色。

只有白黑色时，考虑两种策略。

1. 先选白色或黑色删，最后分开删，答案为树大小一半。
2. 先选白色/黑色叶子删，再删黑色/白色叶子，答案为直径长度一半。

显然策略二更优。

再考虑有灰色点的情况，DP 每个灰点分别选什么色时直径长度最小值，即为答案。

记 $f_{u,1/2}$ 表示将点 $u$ 染成白/黑，当前子树某个点离它最大的距离，$g_{u,1/2}$ 表示将点 $u$ 染成白/黑经过点 $u$ 的答案。

$f$ 的作用是计算直径，不同子树下两最远点到 $u$ 的路径合并成经过 $u$ 的直径。

树的直径就是经过某点且在其子树内的最长路径最大值。

答案就是所有 $u$ 的 $\min(g_{u,1},g_{u,2})$ 的最大值。

考虑转移。

则有

$$f_{u,j}=\max(f_{u,j},\min(f_{v,k}+[j\ne k]))\\ g_{u,j}=\max(g_{u,j},\min(f_{u,j}+f_{v,k}+[j\ne k]))$$

时间复杂度 $\mathcal O(n)$。

#include<bits/stdc++.h>

using namespace std;

//#define int long long
typedef long long ll;

#define ha putchar(' ')
#define he putchar('\n')

inline int read() {
	int x = 0, f = 1;
	char c = getchar();
	while (c < '0' || c > '9') {
		if (c == '-')
			f = -1;
		c = getchar();
	}
	while (c >= '0' && c <= '9')
		x = (x << 3) + (x << 1) + (c ^ 48), c = getchar();
	return x * f;
}

inline void write(int x) {
	if (x < 0) {
		putchar('-');
		x = -x;
	}
	if (x > 9)
		write(x / 10);
	putchar(x % 10 + 48);
}

const int _ = 2e5 + 10, inf = 0x3f3f3f3f;

int n, m, ans, a[_], f[_][2], g[_][3];

vector<int> d[_];

void dfs(int u, int fa)
{
	if(a[u])
	{
		f[u][a[u] ^ 3] = g[u][a[u] ^ 3] = inf;
		f[u][a[u]] = 0, g[u][a[u]] = -inf;
	}
	else
	{
		f[u][1] = f[u][2] = 0;
		g[u][1] = g[u][2] = -inf;
	}
	for(int v : d[u])
	{
		if(v == fa) continue;
		dfs(v, u);
		for(int j = 1; j <= 2; ++j)
		{
			if(a[u] + j == 3) continue;
			int nf = inf, ng = inf;
			for(int k = 1; k <= 2; ++k)
			{
				if(f[v][k] == inf) continue;
				nf = min(nf, f[v][k] + (j != k));
				ng = min(ng, f[u][j] + f[v][k] + (j != k));
			}
			f[u][j] = max(f[u][j], nf), g[u][j] = max(g[u][j], ng);
		}
	}
	ans = max(ans, min(g[u][1], g[u][2]));
}

signed main()
{
	int T = read();
	while(T--)
	{
		n = read();
		for(int i = 1; i <= n; ++i) d[i].clear(), a[i] = read();
		for(int i = 1, u, v; i < n; ++i)
		{
			u = read(), v = read();
			d[u].push_back(v), d[v].push_back(u);
		}
		ans = 0;
		dfs(1, 0);
		write((ans + 1) / 2 + 1), he;
	}
	return 0;
}