A. Timofey and a tree
题意:给一棵树,要求判断是否存在一个点,删除这个点后,所有连通块内颜色一样。$N,C \le 10^5$
想法:这个叫换根吧。先求出一个点合法即其儿子的子树内颜色一样,非该点子树的点颜色都一样。可以用DFS序解决。
#include< cstdio >
typedef long long ll;
template
inline void read(T&x)
{
x=0;bool f=0;char c=getchar();
while((c<'0'||c>'9')&&c!='-') c=getchar();if(c=='-')f=1, c=getchar();
while(c>='0'&&c<='9'){x=x*10+c-'0'; c=getchar();}
x=f?-x:x;
}
const int MAXN(100010);
struct Node{int nd,nx;}bot[MAXN<<1];int tot,first[MAXN];
void add(int a,int b){bot[++tot]=(Node){b,first[a]};first[a]=tot;}
int n,a,b,col[MAXN],Eu[MAXN],size[MAXN],num[MAXN],last[MAXN],cnt,Ans;
bool dif[MAXN];
void DFS(int x,int f)
{
Eu[x]=++cnt; num[cnt]=x;size[x]=1;
for(int v=first[x];v;v=bot[v].nx)
if(bot[v].nd!=f)
{
DFS(bot[v].nd,x); size[x]+=size[bot[v].nd];
dif[x]|=dif[bot[v].nd]|(col[x]!=col[bot[v].nd]);
}
}
void DP(int x,int f)
{
bool bf=false;
for(int v=first[x];v;v=bot[v].nx)
if(bot[v].nd!=f)
{
DP(bot[v].nd,x);
bf|=dif[bot[v].nd];
}
if(!bf)
{
// fprintf(stderr,"%d\n",x);
int L=Eu[x]-1,R=Eu[x]+size[x];
// fprintf(stderr,"%d %d\n",L,R);
if((!L||last[1]>=L)&&(R>n||last[R]>=n)&&(!L||R>n||col[num[1]]==col[num[n]]))Ans=x;
}
}
int main()
{
// freopen("C.in","r",stdin);
read(n);
for(int i=1;i<n;i++)
{
read(a);read(b);
add(a,b);add(b,a);
}
for(int i=1;i<=n;i++)read(col[i]);
DFS(1,0);
for(int i=n;i>=1;i--)last[i]=(col[num[i]]==col[num[i+1]])?last[i+1]:i;
DP(1,0);
if(!Ans)printf("NO");
else printf("YES\n%d\n",Ans);
return 0;
}