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;
}