func_code_string: #include <bits/stdc++.h>   using namespace std;   int cnt[3], ans;   vector<int> c;   vector<vector<int>> sum;   vector<vector<int>> arr;   void dfs(int x, int p) {     for (auto i : arr[x]) {       if (i != p) {         dfs(i, x);         sum[x][0] += sum[i][0];         sum[x][1] += sum[i][1];         sum[x][2] += sum[i][2];       }     }     sum[x][c[x]]++;     if (sum[x][1] == 0 && sum[x][2] == cnt[2] ||         sum[x][1] == cnt[1] && sum[x][2] == 0)       ans++;   }   int main() {     cin.tie(0), ios::sync_with_stdio(0);     int n;     cin >> n;     c.resize(n);     arr.resize(n);     for (int i = 0; i < n; i++) cin >> c[i], cnt[c[i]]++;     for (int i = 0; i < n - 1; i++) {       int a, b;       cin >> a >> b;       arr[a - 1].push_back(b - 1);       arr[b - 1].push_back(a - 1);     }     sum.resize(n, vector<int>(3, 0));     dfs(0, 0);     cout << ans <<   n ;     return 0;   }   