求助!为什么会TLE(从#7开始)
查看原帖
求助!为什么会TLE(从#7开始)
162025
demonlover楼主2023/8/16 16:57
#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
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)
{
    dis[x]=dis[fx]+1;
	for(int i=hd[x];i;i=e[i].next)
	{
		int v=e[i].to;
		if(v==fx) continue;
		dfs1(v,x);
	}
}

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

int main()
{
	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);
	}
	for(int i=1;i<=n;i++)
	{
		for(int j=hd2[i];j;j=e2[j].next)
		{
			if(k[i]!=k[e2[j].to]) 
			{
				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<=n;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/2);
	return 0;
}
2023/8/16 16:57
加载中...