763 A Codeforces - Timofey and a tree
Solution in c++
#include <bits/stdc++.h>
using namespace std;
int n ,t1,t2,g,r=1,rrr,flag,a1,a2;
bool w=0;
vector<int> colors(100005);
map<int,int> in;
vector<pair<int,int> > all;
bool tr(int rr)
{
for(auto u : all)
{
if(colors[u.first]!=colors[u.second])
{
if(u.first!=rr && u.second!=rr)return 0;
}
}
w=1;
rrr=rr;
return 1;
}
int main() {
ios::sync_with_stdio(0);
cin.tie(0);cout.tie(0);
cin>>n;
for(int i=0;i<n-1;i++)
{
cin>>t1>>t2;
all.push_back(make_pair(t1,t2));
}
for(int i=1;i<=n;i++)cin>>colors[i];
for(auto u : all)
{
if(colors[u.first]!=colors[u.second])
{
g++;
if(!flag)
{
flag=1;
a1=u.first;
a2=u.second;
}
in[u.first]++;
in[u.second]++;
if(in[u.first]>1)
{
r=u.first;break;
}
else if(in[u.second]>1)
{
r=u.second;break;
}
}
}
if(g==1)
{
tr(a1);
tr(a2);
}
else {w=tr(r);rrr=r;}
if(w)cout<<"YES\n"<<rrr;
else cout<<"NO";
return 0;
}
0 التعليقات:
إرسال تعليق