func_code_string: #include <bits/stdc++.h>   using namespace std;   const int N = 3e5 + 15;   vector<int> g[N];   int n, m, b[N], r[N], a[N], res;   void dfs(int v, int pa) {     if (a[v] == 1)       r[v]++;     else if (a[v] == 2)       b[v]++;     for (int i : g[v]) {       if (i == pa) continue;       dfs(i, v);       b[v] += b[i];       r[v] += r[i];     }   }   int main() {     ios::sync_with_stdio(false);     cin >> n;     for (int i = 0; i < n; i++) cin >> a[i];     for (int i = 0; i < n - 1; i++) {       int x, y;       cin >> x >> y;       x--, y--;       g[x].push_back(y);       g[y].push_back(x);     }     dfs(0, -1);     for (int i = 1; i < n; i++) {       if (b[i] && r[i]) continue;       int blue = b[0] - b[i], red = r[0] - r[i];       if (((!blue) && red) || (blue && (!red))) {         res++;       }     }     cout << res << endl;   }   