func_code_string: #include <bits/stdc++.h>   using namespace std;   int chk[300004][3];   int visit[300004];   vector<int> vec[300004];   vector<pair<int, int>> vecp;   void dfs(int x) {     visit[x] = 1;     for (int i = 0; i < vec[x].size(); i++) {       if (visit[vec[x][i]] == 0) {         dfs(vec[x][i]);         chk[x][0] += chk[vec[x][i]][0];         chk[x][1] += chk[vec[x][i]][1];         chk[x][2] += chk[vec[x][i]][2];         vecp.push_back({chk[vec[x][i]][1], chk[vec[x][i]][2]});       }     }   }   int main() {     int a, b, c, d, e, m = 0;     cin >> a;     int r = 0, bl = 0;     for (int i = 1; i <= a; i++) {       cin >> b;       if (b == 1)         r++;       else if (b == 2)         bl++;       chk[i][b]++;     }     for (int i = 1; i < a; i++) {       cin >> b >> c;       vec[b].push_back(c);       vec[c].push_back(b);     }     dfs(1);     for (int i = 0; i < vecp.size(); i++) {       int red = vecp[i].first;       int blue = vecp[i].second;       int redo = r - red;       int blueo = bl - blue;       if (red >= 1 && blue == 0 && blueo >= 1 && redo == 0) {         m++;       } else if (red == 0 && blue >= 1 && blueo == 0 && redo >= 1) {         m++;       }     }     cout << m << endl;   }   