Codeforces Round #395 C. Timofey and a tree

Posted 2020-09-01

package codeforces;
import java.util.*;
public class CodeForces_764C_Timofey_and_a_tree {
       static final int N=(int) (2e5+10);
       @SuppressWarnings("unchecked")     
       static ArrayList<Integer> a[]=new ArrayList[N];
       static int book[]=new int[N];
       static int c[]=new int[N];
       static void dfs(int u,int fa)
       {
           int v;
           for(int i=0; i<a[u].size(); i++)
           {
               v=(int) a[u].get(i);
               if(v!=fa)
               {
                   dfs(v,u);
                   if(c[v]!=c[u])
                   {
                       book[u]++;
                       book[v]++;
                   }
               }
           }
       }
       static void solve(){
           Scanner sc=new Scanner(System.in);
           while(sc.hasNext()){
              
               int n = sc.nextInt();
               int u,v;
               for(int i=1; i<n; i++)
               {
                   u=sc.nextInt();
                   v=sc.nextInt();
                   if(a[u]==null) a[u]=new ArrayList<Integer>();
                   if(a[v]==null) a[v]=new ArrayList<Integer>();
                   a[u].add(v);
                   a[v].add(u);
               }
               Arrays.fill(book,0);            
               for(int i=1; i<=n; i++) c[i]=sc.nextInt();
               sc.close();    
               dfs(1,-1);
               int ans2=n,flag=0,maxx=0;
               for(int i=1; i<=n; i++)
               {
                   if(book[i]>maxx)
                   {
                       maxx=book[i];
                       ans2=i;
                   }
               }
               for(int i=0; i<a[ans2].size(); i++)
               {
                   v=(int)a[ans2].get(i);
                   book[ans2]-=book[v];
                   book[v]=0;
               }
               if(book[ans2]!=0) System.out.println("NO");
               else
               {
                   for(int i=1; i<=n; i++)
                   {
                       if(book[i]!=0)
                       {
                           flag=1;
                           break;
                       }
                   }
                   if(flag==1) System.out.println("NO");
                   else
                       System.out.println("YES\n"+ans2);
               }
           }
       }
       public static void main(String args[]){
           solve();
       }
}