func_code_string: #include <bits/stdc++.h>   using namespace std;   long long n, clr[300001][3] = {{0}};   vector<long long> v[300001];   void dfs(long long node, long long parent) {     for (auto x : v[node]) {       if (x == parent) continue;       dfs(x, node);       clr[node][1] += clr[x][1];       clr[node][2] += clr[x][2];     }   }   signed main() {     cin.tie(0);     cout.tie(0);     ios_base::sync_with_stdio(false);     cin >> n;     long long r = 0, b = 0;     for (long long i = 1; i < n + 1; i++) {       long long c;       cin >> c;       if (c == 1) r++;       if (c == 2) b++;       clr[i][c]++;     }     for (long long i = 1; i < n; i++) {       long long x, y;       cin >> x >> y;       v[x].push_back(y);       v[y].push_back(x);     }     dfs(1, 1);     long long ans = 0;     if (b == 0 || r == 0) return cout << n - 1 <<   n , 0;     for (long long i = 1; i < n + 1; i++) {       if ((clr[i][2] == 0 && clr[i][1] == r) ||           (clr[i][1] == 0 && clr[i][2] == b))         ans++;     }     cout << ans;   }   