【CF734E】Anton and Tree【树的直径】

分析

一开始理解错题意了，直接暴力判断连通块个数，我就说怎么这么简单。。。

其实这个颜色的连通块是会动态变化的，染色之后可能会有合并连通块的发生。问题好像瞬间复杂了起来。

但是在染色的时候连通块一定会一起被染色，于是就相当于一个点了，我们也用DFS去缩成一个点，不难发现，新的一棵树上，都是黑白点交替出现的。想到这里可能觉得跟树的深度有关，但是换根深度也会变化。

实际上答案就是(树的直径+1)/2，从直径的中点开始染色就行，只要想到直径方案是显然的。

注意：建新图的时候要判重，不然会多两条边，然后所有子树的搜索次数都会指数增加。要不直接建单向边，然后对面那个连回来。血的教训！！

上代码

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<set>
using namespace std;

struct node
{
	int to,next;
}e[400010],e2[400010];

int n,col[200010],dis[200010],dis2[200010];
int k[200010],cnt,vis[200010];
int tot,hd[400010];
int tot2,hd2[400010];

int read() 
{
    int x=0;char c=getchar();
    while(c<'0'||c>'9') c=getchar();
    while(c>='0'&&c<='9') 
    {
	    x=x*10+c-'0';
		c=getchar();
	}
    return x;
}

void add(int x,int y)
{
	e[++tot]=(node){y,hd[x]};
	hd[x]=tot;
}
void add2(int x,int y)
{
	e2[++tot2]=(node){y,hd2[x]};
	hd2[x]=tot2;
}

void dfs(int x)
{
	vis[x]=1;
	k[x]=cnt;
	for(int i=hd2[x];i;i=e2[i].next)
	{
		int v=e2[i].to;
		if(vis[v]||col[v]!=col[x]) continue;
		vis[v]=1;
		dfs(v);
	}
}

void dfs1(int x,int fx)
{
	for(int i=hd[x];i;i=e[i].next)
	{
		int v=e[i].to;
		if(v==fx) continue;
		dis[v]=dis[x]+1;
		dfs1(v,x);
	}
}

void dfs2(int x,int fx)
{
	for(int i=hd[x];i;i=e[i].next)
	{
		int v=e[i].to;
		if(v==fx) continue;
		dis2[v]=dis2[x]+1;
		dfs2(v,x);
	}
}

int main()
{
//	freopen("a.in","r",stdin);
//	freopen("a.out","w",stdout);
	n=read();
	for(int i=1;i<=n;i++) col[i]=read();
	for(int i=1;i<=n-1;i++)
	{
		int x,y;
		x=read();y=read();
		add2(x,y);
		add2(y,x);
	}
	for(int i=1;i<=n;i++)
	{
		if(!vis[i]) cnt++,dfs(i);
	}
	set<pair<int,int> > S;
	for(int i=1;i<=n;i++)
	{
		for(int j=hd2[i];j;j=e2[j].next)
		{
			if(k[i]!=k[e2[j].to]&&!S.count(make_pair(k[i],k[e2[j].to]))) 
			{
				S.insert(make_pair(k[i],k[e2[j].to]));//记录，判重
				S.insert(make_pair(k[e2[j].to],k[i]));
				add(k[i],k[e2[j].to]);
				add(k[e2[j].to],k[i]);
			}
		}
	}
	dfs1(1,0);
	int mx=0,mxp;
    for(int i=1;i<=cnt;i++)
    {
    	if(dis[i]>=mx) 
    	{
    		mx=dis[i];
    		mxp=i;
		}
	}
	dfs2(mxp,0);
	mx=0;
	for(int i=1;i<=cnt;i++) mx=max(mx,dis2[i]);
	printf("%d",(mx+1)/2);
	return 0;
}