Codeforces764C【DFS】

前言，根据最终图的样貌搞真厉害

“缩点判根度为结点数-1”牛逼

-----

题意：

找一个根使得不带根的所有子树内部颜色都相同；

思路：

如果存在两个颜色不一样的连在一起，根就是他们两个中间一个。。。。。。直接搜两次就好了。。。

#include<bits/stdc++.h>
using namespace::std;
typedef pair<int,int> PII;

const int N = 4e5+10;

struct Edge{
    int to;
    int next;
};
Edge e[N];
int head[N],tol,n,col[N];
bool vis[N];

void add(int u,int v)
{
    e[tol].to=v;
    e[tol].next=head[u];
    head[u]=tol++;
}

void init()
{
    tol=0;
    memset(head,-1,sizeof(head));
}
bool ans;
bool flat;
void DFS(int u,int color)
{
    vis[u]=true;
    if(!ans) return;
    for(int i=head[u];~i;i=e[i].next)
    {
        int v=e[i].to;
        if(vis[v])
            continue;
        vis[v]=true;
        if(col[v]!=color){
            ans=false;
            return;
        }
        DFS(v,color);
    }
}

bool Judge(int x)
{
    ans=true;
    for(int i=1;i<=n;i++) vis[i]=false;
    vis[x]=true;
    for(int i=head[x];~i;i=e[i].next)
    {
        int v=e[i].to;
        DFS(v,col[v]);
        if(!ans)
            return false;
    }
    return true;
}

int solve(int x)
{
    int v;
    bool flag=false;
    for(int i=head[x];~i;i=e[i].next)
    {
        v=e[i].to;
        if(col[v]!=col[x])
        {
            flag=true;
            return v;
        }
    }
    return -1;
}

int main()
{
    int u,v;
    scanf("%d",&n);
    init();
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&u,&v);
        add(u,v);
        add(v,u);
    }
    for(int i=1;i<=n;i++)
        scanf("%d",&col[i]);
    flat=false;
    for(int i=1;i<=n;i++)
    {
        int j;
        j=solve(i);
        if(j!=-1)
        {
            if(Judge(j))
            {
                puts("YES");
                printf("%d\n",j);
            }
            else if(Judge(i))
            {
                puts("YES");
                printf("%d\n",i);
            }
            else
                puts("NO");
            return 0;
        }
    }
    puts("YES");
    puts("1");
    return 0;
}