CodeForces - 1118F1 Tree Cutting (Easy Version)（DFS）

💇‍♀️ 💇‍♀️ 💇‍♀️

//目的：找到一个这样的结点，他的子树包含所有的蓝色/红色，且不包含另一种颜色
//所以记录下每个节点的子树的蓝色，红蓝结点分别有多少
//然后枚举每条边检验
int n,ans;
int a,b;//一共有多少蓝/红色
int tag[MAXN];//某点的颜色
int col[3][MAXN];//某点的子树的颜色个数
int dep[MAXN];//深度
vector<int>mp[MAXN];
void dfs(int x,int fx)
{
    ++col[tag[x]][x];
    if(dep[x]==0) dep[x]=dep[fx]+1;
    for(auto v:mp[x])
    {
        if(v!=fx)
        {
            dfs(v,x);
            col[2][x]+=col[2][v];
            col[1][x]+=col[1][v];
        }
    }
}
void solve()
{
    int n;cin>>n;
    rpp(i,n) 
    {
        cin>>tag[i];
        if(tag[i]==1) ++a;
        if(tag[i]==2) ++b;
    }
    vector<pii>road(n-1);
    rep(i,n-1)
    {
        int x,y;cin>>x>>y;
        road.push_back(make_pair(x,y));
        mp[x].push_back(y);
        mp[y].push_back(x);
    }
    dep[0]=0;
    dfs(1,0);
    for(auto x:road)
    {
        if(dep[x.first]<dep[x.second]) swap(x.first,x.second);
        if(col[1][x.first]==a&&col[2][x.first]==0) ++ans;
        else if(col[2][x.first]==b&&col[1][x.first]==0) ++ans;
    }
    cout<<ans<<endl;
}
signed main()
{
    solve();
    stop;
    return 0;
}