codeforces 734E 并查集缩点 + 树的直径

题目连接：https://codeforc.es/contest/734/problem/E

题意

  给出一颗生成树（边为n-1），树的每一点有两种颜色，现有一种操作：能将与某一点同一种颜色的点集颜色翻转，求至少需要多少次这样的操作能将整棵树变为同一种颜色。

思路

  并查集缩点，用两个集合分别存初始黑白两色的集合，同一种颜色的集合看成一个点，将不同颜色的缩点进行重建图。在新图上只需要求出树的直径d，因为是黑白交替，所以最少的操作就是（d+1）/2

#include<bits/stdc++.h>
using namespace std;
#define maxn 400005
#define maxm 100005

struct edge
{
    int next,to,len;
}G[2*maxn];
int head[maxn],num;
void add(int from,int to,int len)
{
    G[++num].next=head[from];
    G[num].to=to;
    G[num].len=len;
    head[from]=num;
}

int color[maxn];
int ff1[maxn],ff2[maxn];
int find1(int x)
{
    return ff1[x]==x?x:ff1[x]=find1(ff1[x]);
}
int find2(int x)
{
    return ff2[x]==x?x:ff2[x]=find2(ff2[x]);
}

struct node
{
    int u,v;
};

int dis[maxn],Max,t;
bool vis[maxn];
int bfs(int st)
{
    Max=0,t=st;
    memset(dis,0,sizeof(dis));
    memset(vis,false,sizeof(vis));
    queue<int>q;
    q.push(st);
    while(!q.empty())
    {
        int u=q.front();
        vis[u]=true;
        q.pop();
        for(int i=head[u];i!=-1;i=G[i].next)
        {
            int v=G[i].to,len=G[i].len;
            if(!vis[v])
            {
                dis[v]=dis[u]+len;
                q.push(v);
                if(dis[v]>Max)
                {
                    Max=dis[v];
                    t=v;
                }
            }
        }
    }
}

int solve(int st)
{
    bfs(st);
    bfs(t);
    return Max;
}

int main()
{
    vector<node>ff;
    memset(head,-1,sizeof(head));
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&color[i]);
        ff1[i]=i,ff2[i]=i;
    }

    for(int i=0;i<n-1;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        ff.push_back(node{x,y});
        if(color[x]==0&&color[y]==0)
        {
            int nx=find1(x),ny=find1(y);
            if(nx!=ny)
                ff1[nx]=ny;
        }
        if(color[x]&&color[y])
        {
            int nx=find2(x),ny=find2(y);
            if(nx!=ny)
                ff2[nx]=ny;
        }
    }

    for(int i=1;i<=n;i++)
        find1(i),find2(i);

    int st=-1;
    for(int i=0;i<n-1;i++)
    {
        int x=ff[i].u,y=ff[i].v;
        if(color[x]!=color[y])
        {
            if(color[x])
                swap(x,y);
            add(ff1[x],ff2[y]+n,1);
            add(ff2[y]+n,ff1[x],1);
            st=ff1[x];
        }
    }
    if(st==-1)
        printf("0");
    else
    {
        solve(st);
        printf("%d",(Max+1)/2);
    }
    return 0;
}