[Luogu] CF763A Timofey and a tree

https://www.luogu.com.cn/problem/CF763A

前言

写了我一个小时 卧槽 卡死了我了
一开始以为是 并查集缩点然后 dfs枚举
结果WA了 搞得头疼 然后看了一下题解 ？ ？ ？ ？ 直接枚举就行？

思路

• 先处理 出来 有多少个分支 (即两个结点的颜色不同)
• 然后再利用 dfs枚举每一个结点作为 根结点
• 利用分支数相同这个 特点进判断

虽然思路很简单 就一个枚举
但是总感觉卡时间复杂度了 n =1e5跑一遍全图 也就是n*(n-1)/2 (emm)

还有就是dfs 好久没写了 感觉还挺难emm

CODE:

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+10;
int cnt;
int h[N],e[N],ne[N],idx,c[N];
int n,a,b;
void add(int a,int b)
{
    e[idx] = b,ne[idx] = h[a],h[a]=idx++;
}

void dfs(int x,int b)
{
    for(int i = h[x]; i!=-1; i=ne[i])
    {
        int j = e[i];
        if(j == b) continue;
        if(c[x]!=c[j])cnt++;

        dfs(j,x);
    }
}

void solve()
{
    cin>>n;
    memset(h,-1,sizeof h);

    for(int i = 1; i<=n-1; i++)
    {
        cin>>a>>b;
        add(a,b);
        add(b,a);
    }
    for(int i=1; i<=n; i++)
        cin>>c[i];

    dfs(1,0);
    for(int i=1; i<=n; i++)
    {
        int sum  = 0;
        for(int j = h[i]; j!=-1; j=ne[j])
        {
         //   cout<<"debug"<<endl;
            int p  = e[j];
            if(c[p]!=c[i])sum++;

        }
        if(sum == cnt)
        {
            cout<<"YES"<<endl;
          //  puts("YES");
            cout<<i<<endl;
            return;
        }
    }
    cout<<"NO"<<endl;
   //puts("NO");

}
int main()
{
    ios::sync_with_stdio(false);
    solve();
    return 0;
}