P10725 [GESP202406 八级] 最远点对

P10725 [GESP202406 八级] 最远点对 - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

 

题目描述

小杨有⼀棵包含 nn 个节点的树，这棵树上的任意⼀个节点要么是白色，要么是黑色。

小杨想知道相距最远的一对不同颜色节点的距离是多少。

输入格式

第一行包含⼀个正整数 nn，代表树的节点数。

第二行包含 nn 个非负整数 a1,a2,⋯ ,an（对于所有的 1≤i≤n1≤i≤n，均有 ai​ 等于 00 或 11），其中如果 ai=0ai​=0，则节点 ii 的颜色为白色；如果 ai=1 , 则节点 ii 的颜色为黑色。

之后 (n−1)(n−1) 行，每行包含两个正整数 xi,yixi​,yi​，代表存在一条连接节点 xixi​ 和 yiyi​ 的边。

保证输入的树中存在不同颜色的点。

输出格式

输出⼀个整数，代表相距最远的一对不同颜色节点的距离。

输入输出样例

输入 #1复制

5
0 1 0 1 0
1 2
1 3
3 4
3 5

输出 #1复制

3

说明/提示

样例解释

相距最远的不同颜色的一对节点为节点 22 和 55。

数据范围

子任务编号	得分	nn	aiai​	特殊条件
11	3030	≤105≤105	0≤ai≤10≤ai​≤1	树的形态为一条链
22	3030	≤103≤103	0≤ai≤10≤ai​≤1
33	4040	≤105≤105	0≤ai≤10≤ai​≤1

对于全部数据，保证有 1≤n≤1051≤n≤105，0≤ai≤10≤ai​≤1。

CODE

#include<bits/stdc++.h>
#define int long long
using namespace std;

int col[100010],d[100010],ans;
vector<int>g[100010];
void dfs(int u,int f){
	for(auto v:g[u]){
		if(v==f)continue;
		d[v]=d[u]+1;
		dfs(v,u);
	}
}
void get_ans(int u,int f,int t,int l){
	if(col[u]==t)ans=max(ans,l);
	for(auto v:g[u]){
		if(v==f)continue;
		get_ans(v,u,t,l+1);
	}
}
signed main(){	
	int n;
	cin>>n;
	for(int i=1;i<=n;i++)cin>>col[i];
	for(int i=1;i<n;i++){
		int u,v;
		cin>>u>>v;
		g[u].push_back(v);
		g[v].push_back(u);
	}
	dfs(1,0);
	int id0=-1,id1=-1;
	int mx0=-1,mx1=-1;
	for(int i=1;i<=n;i++){
		if(col[i]==0){
			if(d[i]>mx0){
				mx0=d[i];
				id0=i;
			}
		}else{
			if(d[i]>mx1){
				mx1=d[i];
				id1=i;
			}
		}
	}
	if(id0!=-1)get_ans(id0,0,1,0);
	if(id1!=-1)get_ans(id1,0,0,0);
	cout<<ans;
	return 0;
}