F1. Tree Cutting (Easy Version)(思维+树dfs)

https://codeforces.com/problemset/problem/1118/F1

题意翻译

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

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

问有多少条合法的边。

数据范围：1 <= n <= 3e5，a_iai​ = 1表示是红色，a_iai​ = 2 表示蓝色

输入输出样例

输入 #1复制

5
2 0 0 1 2
1 2
2 3
2 4
2 5

输出 #1复制

1

输入 #2复制

5
1 0 0 0 2
1 2
2 3
3 4
4 5

输出 #2复制

4

输入 #3复制

3
1 1 2
2 3
1 3

输出 #3复制

0

---

思路：最开始是暴力枚举每一条被去掉的边再每次dfs一次，O(n^2）

正解：统计每个子树内红色和蓝色的数量，当子树内只有一种颜色兵器这种颜色等于这种颜色总数的时候，那么子树根顶连的那条就能++（此时满足整张图的去掉子树的部分不会有该种颜色出现)

#include<iostream>
#include<vector>
#include<queue>
#include<cstring>
#include<cmath>
#include<map>
#include<set>
#include<cstdio>
#include<algorithm>
#define debug(a) cout<<#a<<"="<<a<<endl;
using namespace std;
const int maxn=3e5+100;
typedef long long LL;
LL a[maxn],col[maxn][3],ans=0;
vector<LL>g[maxn];
LL sum1,sum2;
void dfs(LL u,LL fa)
{
	if(a[u]==1) col[u][1]++;
	if(a[u]==2) col[u][2]++;
	for(LL i=0;i<g[u].size();i++){
		LL v=g[u][i];
		if(v==fa) continue;
		dfs(v,u);
		col[u][1]+=col[v][1];
		col[u][2]+=col[v][2];
	}
	if((col[u][1]==sum1&&col[u][2]==0)||(col[u][1]==0&&col[u][2]==sum2))
	{
		ans++;
	}
}
int main(void)
{
  cin.tie(0);std::ios::sync_with_stdio(false);
  LL n;cin>>n;
  for(LL i=1;i<=n;i++){
  	cin>>a[i];
  	if(a[i]==1) sum1++;
  	if(a[i]==2) sum2++;
  } 
  for(LL i=1;i<n;i++){
  	LL x,y;cin>>x>>y;
  	g[x].push_back(y);g[y].push_back(x);
  }
  dfs(1,0);
  cout<<ans<<endl;
return 0;
}