大体题意:
给你一棵树,要求找一个结点来作为根节点,使得所有的子树颜色都一样,整棵树不是子树,其余的是,存在的话输出那个点。
思路:
并查集+缩点。
先把颜色一样的,并且连在一起的点缩成一个点。(类似连通块) 缩点后也一定是树。
这样最后的形式就是 一个根结点连接着很多其他结点,这样枚举所有的边,统计一下缩点的度数。
找出那个度数为总和-1的那个点就是答案,没有的话,就是NO咯。
#include <bits/stdc++.h>
#define ps push_back
#define fi first
#define se second
#define mr make_pair
using namespace std;
const int inf = 0x3f3f3f3f;
const double eps = 1e-10;
const double pi = acos(-1.0);
typedef long long ll;
typedef unsigned long long ULL;
const int maxn = 100000 + 10;
vector<int>g[maxn];
int c[maxn];
int d[maxn];
int n;
struct Node{
int f,t;
void read(){
scanf("%d %d",&f, &t);
}
}p[maxn];
int de[maxn];
int fa[maxn];
int cnt = 0;
void dfs(int cur,int la){ ///dfs 进行缩点。
de[cur] = 1;
fa[cur] = cnt;
for (int i = 0; i < g[cur].size(); ++i){
int v = g[cur][i];
if (v!=la && c[v]==c[cur]){
dfs(v,cur);
}
}
}
int main(){
scanf("%d",&n);
memset(fa,-1,sizeof fa);
for (int i = 0; i < n-1; ++i){
p[i].read();/// 记录边
int u = p[i].f;
int v = p[i].t;
g[u].ps(v);/// 便于dfs
g[v].ps(u);
}
for (int i = 1; i <= n; ++i) {
scanf("%d",c+i);
}
for (int i= 1; i <= n; ++i){
if (!de[i]){/// 已经加入连通块的点 就不在dfs了。
cnt++;
dfs(i,-1);
}
}
memset(de,0,sizeof de);
bool ok = 1;
for (int i = 0; i < n-1; ++i){
int u = p[i].f;
int v = p[i].t;
int fu = fa[u];
int fv = fa[v];
if (fa[u] != fa[v]){
d[u]++;///统计度数
d[v]++;
}
}
for (int i = 1 ; i <= n; ++i){
if (d[i] == cnt-1) return 0 * printf("YES\n%d\n",i);
}
puts("NO");///这里的NO 是找不到那个根结点,度数为所有点-1.
return 0;
}