Codeforces Round #395 C

Timofey and a tree

题意：给一颗树，每个节点有一个颜色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;
}