Codeforces 734E 树的直径+缩点

点击打开链接

/颜色相同的联通块能同时变色,则缩成一个点

//缩点后 变成一个黑白相间的树 
//选中心点变色使得同色的一圈向外扩散 最多扩散(d+1)/2次 d为树的直径 

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <map>
#include <vector>
using namespace std;
typedef long long ll;
const int N=2e5+20;
int n,a[N],vis[N];
vector<int> e[N],g[N];
int num=0,bel[N];
pair<int,int> p[N];
void dfs(int u)
{
	vis[u]=1;
	bel[u]=num;
	for(int i=0;i<e[u].size();i++)
	{
		int v=e[u][i];
		if(!vis[v]&&a[u]==a[v])//颜色相同的联通块缩成一个点
		dfs(v);
	}
}
int dp[N][2];//dp[u][0 or 1] 经过u的最长和次长叶节点 
void dfs(int u,int fa)
{
	dp[u][0]=dp[u][1]=0;
	for(int i=0;i<g[u].size();i++)
	{
		int v=g[u][i];
		if(v==fa) continue;
		
		dfs(v,u);
		if(dp[v][0]+1>dp[u][0])
		{
			dp[u][1]=dp[u][0];
			dp[u][0]=dp[v][0]+1;
		}
		else if(dp[v][0]+1>dp[u][1])
		dp[u][1]=dp[v][0]+1; 
	}
}
int main()
{
	memset(vis,0,sizeof(vis));
	cin>>n;
	for(int i=1;i<=n;i++)
	scanf("%d",&a[i]);
	
	for(int i=1;i<=n-1;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		e[u].push_back(v);
		e[v].push_back(u);
	
		p[i].first=u;
		p[i].second=v;
	}
	
	for(int i=1;i<=n;i++)
	{
		if(!vis[i])//颜色相同的联通块缩成一个点 
		{
			num++;
			dfs(i);
		}
	}
	for(int i=1;i<n;i++)
	{
		int u=p[i].first;
		int v=p[i].second;
		if(bel[u]!=bel[v])
		{
			g[bel[u]].push_back(bel[v]);
			g[bel[v]].push_back(bel[u]);
		}
	}
	//缩点后 变成一个黑白相间的树 
	//选中心点变色使得同色的一圈向外扩散 最多扩散(d+1)/2次 d为树的直径 
	dfs(1,0);//
	int d=0;
	for(int i=1;i<=n;i++)
	{
		d=max(d,dp[i][0]+dp[i][1]);
	}
	cout<<(d+1)/2<<endl;
	
	return 0;
}