Codeforces Round #379 (Div. 2) E. Anton and Tree

版权声明：原创文章，转载要注明作者哦 https://blog.csdn.net/DYT_B/article/details/81433454

题目描述：戳这里
题解：
这题还是比较妙的。
我们发现有这么多的同色块，块的大小又不一样，好像无从下手。
但是我们发现一个块最多只要改一次，那么我们可以把每一个块缩成一个点（dfs一下）就省去了很多麻烦。
缩完点之后，题目就变得简单了很多，一个黑点边上只有白点，白点边上只有黑点。这就可以贪心解决，答案肯定是树的直径的一半。

代码如下：

#include<cstdio>
#include<string>
#include<cstring>
using namespace std;
const int maxn=200005;
int n,m,tot,sum,col[maxn],lnk[maxn],son[2*maxn],nxt[2*maxn],fa[maxn],s[maxn],dis[2][maxn];
struct dyt{
    int x,y;
}a[maxn];
void add(int x,int y){
    son[++tot]=y,nxt[tot]=lnk[x],lnk[x]=tot;
}
void clear(){
    tot=0;
    memset(lnk,0,sizeof(lnk));
    memset(nxt,0,sizeof(nxt));
}
void dfs(int x,int num){
    fa[x]=num;
    for (int j=lnk[x];j;j=nxt[j])
    if (col[son[j]]==col[x]&&!fa[son[j]]) dfs(son[j],num);
}
void dfs1(int x,int fat,int num){
    for (int j=lnk[x];j;j=nxt[j])
    if (son[j]!=fat) {
        dis[num][son[j]]=dis[num][x]+1;
        dfs1(son[j],x,num);
    }
}
void make_(){
    dfs1(1,0,0); int id=1;
    for (int i=2;i<=m;i++)
    if (dis[0][i]>dis[0][id]) id=i;
    dfs1(id,0,1); int id1=1;
    for (int i=2;i<=m;i++)
    if (i!=id&&dis[1][i]>dis[1][id1]) id1=i;
    printf("%d\n",(dis[1][id1]+1)/2);
}
int main(){
    scanf("%d",&n);
    for (int i=1;i<=n;i++) scanf("%d",&col[i]);
    for (int i=1;i<n;i++) {
        scanf("%d %d",&a[i].x,&a[i].y);
        add(a[i].x,a[i].y); add(a[i].y,a[i].x);
    }
    for (int i=1;i<=n;i++) if (!fa[i]) dfs(i,i),s[i]=++m;
    clear();
    for (int i=1;i<=n;i++)
    if (s[fa[a[i].x]]!=s[fa[a[i].y]]){
        add(s[fa[a[i].x]],s[fa[a[i].y]]);
        add(s[fa[a[i].y]],s[fa[a[i].x]]);
    }
    make_();
    return 0;
}