题意:给一颗树,每个节点有一个颜色c[i],问是否存在一个点,使得去掉这个点后每颗树的颜色只有一种,如果存在,输出这个点
思路:dfs+剪枝做的,如果从结点u->v是可行的(也就是把u删除后v为根的数只有一种颜色)那么标记,之后再次dfs的时候便不需要重新走,如果不可行,也标记起来,同样下次dfs的时候直接返回结果,不需要重新遍历,这样优化后复杂度为O(N) ,A掉之后看了大佬的代码都是直接判断边2端的颜色是不是一样,如果不是一样必然有一个是要删掉的,然后记录颜色不一样的点的度,并且记录这样的点对有多少,删除的点的度数一定是等于这样点对的个数,判断就可以了,当然由于dfs优化的姿势还是要学习一个的,所以我就不贴大佬的代码了,直接贴我自己的dfs吧
AC代码:
#include "iostream" #include "string.h" #include "stack" #include "queue" #include "string" #include "vector" #include "set" #include "map" #include "algorithm" #include "stdio.h" #include "math.h" #define ll long long #define bug(x) cout<<x<<" "<<"UUUUU"<<endl; #define mem(a) memset(a,0,sizeof(a)) #define mp(x,y) make_pair(x,y) #define pb(x) push_back(x) using namespace std; const long long INF = 1e18+1LL; const int inf = 1e9+1e8; const int N=1e5+100; const ll mod=1e9+7; int n,c[N]; vector<int> vex[N]; map<pair<int,int>,int> vis,vis0; map<int,int> M; inline int dfs(int fa, int u, int r){ if(c[u]!=r) return 0; for(auto i : vex[u]){ if(i==fa || vis[mp(u,i)]) continue; if(vis0[mp(u,i)]){ return 0; } if(!dfs(u,i,c[u])){ vis0[mp(u,i)]=1; return 0; } else vis[mp(u,i)]=1; } return 1; } void add(int u, int v){ vex[u].pb(v); vex[v].pb(u); } int main(){ int u,v; ios::sync_with_stdio(false),cin.tie(0),cout.tie(0); cin>>n; for(int i=1; i<n; ++i){ cin>>u>>v; add(u,v); } for(int i=1; i<=n; ++i){ cin>>c[i]; M[c[i]]++; } int m=M.size(); int ans; for(int i=1; i<=n; ++i){ ans=i; if(vex[i].size()>=m-1){ //cout<<i<<endl; for(auto j : vex[i]){ if(vis[mp(i,j)]) continue; if(!dfs(i,j,c[j])){ ans=0; break; } } } else ans=0; if(ans) break; } if(ans){ cout<<"YES"<<endl; cout<<ans<<endl; } else cout<<"NO"; return 0; }