codeforces 1118F1 540 div3

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