洛谷P10725题解

前言

在夏令营的时候学了树的直径，看完了题目描述之后就觉得应该差不多。

思路

用 _1 来存离i最远的黑色节点的距离。

用 _0 来存离i最远的白色节点的距离。

考虑两种情况：

1. 左端点是黑，右端点是白。

2. 左端点是白，右端点是黑。

分别更新一次答案。

遍历完一个点还有更新 _1 和 _0 。

具体见代码注释。

代码

#include <bits/stdc++.h>
using namespace std;
//个人习惯
//#define int long long
vector<int> a[100005];
bool c[100005];
int ans = -1e9;
int _1[100005];//离i最远的黑色节点的距离
int _0[100005];//离i最远的白色节点的距离
void AAA(int x,int fa)
{
//	x是当前节点
//	fa是父节点
	if(c[x]) _1[x] = 0;
	else _0[x] = 0;
	for(int i=0;i<a[x].size();i++)
	{
		int y = a[x][i];
		if(y==fa) continue;
//		不是父节点就往下搜
		AAA(y,x);
//		更新答案
		if(_0[x]>=0&&_1[y]>=0) ans = max(ans,_1[y]+_0[x]+1);
		if(_1[x]>=0&&_0[y]>=0) ans = max(ans,_0[y]+_1[x]+1);
		_1[x] = max(_1[x],_1[y]+1);
		_0[x] = max(_0[x],_0[y]+1);
	}
}
signed main()
{
//	个人习惯
//	freopen(".in","r",stdin);
//	freopen(".out","w",stdout);
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	int n;
	cin>>n;
//	初始化
	for(int i=1;i<=n;i++)
	{
		cin>>c[i];
//		一定要小于-1e5，不然会WA
		_1[i] = -1e9;
		_0[i] = -1e9;
	}
//	输入+建图
	for(int i=1;i<n;i++)
	{
		int u,v;
		cin>>u>>v;
		a[u].push_back(v);
		a[v].push_back(u);
	}
	AAA(1,0);
	cout<<ans;
	return 0;
}