Codeforces CF1118F1 Tree Cutting (Easy Version) 题解

题目简述

给一棵树，每一个节点都是红色，蓝色或者无色。

一条边是合法的当且仅当删除这一条边之后，树被分成两部分，这两部分不同时含有红色和蓝色。

问有多少条合法的边。

题目分析

考虑先搜索一遍这棵树，并处理出以 \(x\) 节点为根的子树中红点和蓝点的数量，记为 \(red_x\) 和 \(blue_x\)。树中所有的蓝点和红点的数量记为 \(\texttt{BLUE}\) 和 \(\texttt{RED}\)。

如果 \(u \rightarrow v\) 这条边是合法的，那么以 \(v\) 为根的子树中仅有一种有颜色的点，即以下两种情况：

• \(red_v=\texttt{RED}\) 并且 \(blue_v=0\)。
• \(blue_v=\texttt{BLUE}\) 并且 \(red_v=0\)。

在搜索的时候判断即可。

代码

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define random(a,b) (rand()%(b-a+1)+a)
const int N=3e5+10;
vector<int> G[N];
int n,a[N],Blue,Red,ans;
pair<int,int> dfs(int u,int father)
{
	int blue=(a[u]==2),red=(a[u]==1);
	for(int v:G[u])
	{
		if(v!=father)
		{
			auto temp=dfs(v,u);
			int A=temp.first,B=temp.second;
			if((A==Blue&&B==0)||(A==0&&B==Red)) ans++;
			blue+=A,red+=B;
		}
	}
	return {blue,red};
}
int main()
{
	ios::sync_with_stdio(false);
    cin.tie(0);
    cin>>n;
    for(int i=1;i<=n;i++)
    {
    	cin>>a[i];
    	if(a[i]==2) Blue++;
    	if(a[i]==1) Red++;
	}
	for(int u=0,v=0,i=1;i<n;i++)
	{
		cin>>u>>v;
		G[u].push_back(v);
		G[v].push_back(u);
	}
	dfs(1,0);
	cout<<ans;
	return 0;
}