CodeForces

直接缩点,然后跑一次树的直径就好了....
至于跑树的直径
简易推导可得,原图的答案其实等价与两个最原点....
因为(得到答案)就像水流一样,...一层一层的...
然后就结束了

#include<bits/stdc++.h>
#define MAXN 200005
using namespace std;

int n,tot,tt,tot2,maxl1=0,maxl2=0,maxl=0;
int h[MAXN],h2[MAXN],color[MAXN];//1 white 2 black 
int belong[MAXN],dis[MAXN];

struct node{
    int from,to,next;
}e[MAXN<<1],e2[MAXN<<1];

void init(){
    tot=tt=tot2=0;
    memset(h,-1,sizeof(h));
    memset(h2,-1,sizeof(h2));
    memset(belong,-1,sizeof(belong));
}

void add(int x,int y){
    tot++;
    e[tot].from=x;
    e[tot].to=y;
    e[tot].next=h[x];
    h[x]=tot;
}

void dfs(int now,int fa){
    if(belong[now]==(-1)){
        tt++;
        belong[now]=tt;
    }
    for(int i=h[now];i!=(-1);i=e[i].next){
        if(e[i].to!=fa){
            if(color[e[i].to]==color[now])belong[e[i].to]=belong[now];
            dfs(e[i].to,now);
        }
    }
}

void add2(int x,int y){
    tot2++;
    e2[tot2].from=x;
    e2[tot2].to=y;
    e2[tot2].next=h2[x];
    h2[x]=tot2;
}

int dfs2(int now,int fa){
    if(dis[now]>dis[maxl]){
        maxl=now;    
    }
    for(int i=h2[now];i!=(-1);i=e2[i].next){
        if(e2[i].to!=fa){
            dis[e2[i].to]=dis[now]+1;
            dfs2(e2[i].to,now);
        }
    }
}

int main(){
    init();
    cin>>n;
    for(int i=1;i<=n;i++)cin>>color[i],color[i]++;
    for(int i=1;i<n;i++){
        int x,y;
        cin>>x>>y;
        add(x,y);
        add(y,x);
    }
    dfs(1,1);
    for(int i=1;i<=tot;i++){
        if(belong[e[i].from]!=belong[e[i].to]){
            add2(belong[e[i].from],belong[e[i].to]);
        }
    }
    //找树的直径 
    memset(dis,0,sizeof(dis));
    dfs2(1,1);
    memset(dis,0,sizeof(dis));
    maxl1=maxl;
    dfs2(maxl1,maxl1);
    cout<<(dis[maxl]+1)/2<<endl;
    
}