codeforces 1118F1
dfs记录每个节点包括其子节点的各个颜色顶点记为C[V][3]
Sum[i] - C[V][i] 及为另一分量i颜色的个数
#include<bits/stdc++.h>
using namespace std;
#define fst first
#define sec second
#define sci(num) scanf("%d",&num)
#define scl(num) scanf("%lld",&num)
#define mem(a,b) memset(a,b,sizeof a)
#define cpy(a,b) memcopy(a,b,sizeof b)
typedef long long LL;
typedef pair<int,int> P;
const int MAX_N = 3e5 + 100;
const int MAX_M = 310;
struct edge{
int to,next;
} es[MAX_N << 1];
int head[MAX_N],len;
int color[MAX_N];
int C[MAX_N][3];
int sum[3];
void add(int from,int to) {
es[len].to = to;
es[len].next = head[from];
head[from] = len++;
}
int ans = 0;
void dfs(int u,int f,int s) {
C[u][color[u]]++;
for (int i = head[u]; i != -1;i = es[i].next) {
edge& e = es[i];
if (e.to == f) continue;
dfs(e.to,u,s|color[u]);
for (int i = 0;i < 3;i++)
C[u][i] += C[e.to][i];
}
int cnt1 = 0;
for (int i = 1;i < 3;i++)
if (C[u][i] != 0) cnt1++;
int cnt2 = 0;
for (int i = 1;i < 3;i++)
if (sum[i] - C[u][i] != 0) cnt2++;
if (cnt1 <= 1 && cnt2 <= 1) ans++;
}
int main() {
ios::sync_with_stdio(false); cin.tie(0);
mem(head,-1); len = 0;
int N;
cin >> N;
for (int i = 1;i <= N;i++) {
cin >> color[i];
sum[color[i]]++;
}
for (int i = 1;i < N;i++) {
int u,v;
cin >> u >> v;
add(u,v);
add(v,u);
}
ans = 0;
dfs(1,-1,0);
cout << ans << endl;
return 0;
}