func_code_string: #include <bits/stdc++.h>   using ll = long long;   using namespace std;   const int MOD = 1e9 + 7;   map<int, vector<int>> nodecolor;   map<int, vector<int>> tree;   void dfs(vector<int>& vis, int node, vector<int>& color) {     int blue = 0, c = 0, red = 0;     if (color[node - 1] == 0) {       c += 1;     }     if (color[node - 1] == 1) {       red += 1;     }     if (color[node - 1] == 2) {       blue += 1;     }     for (auto x : tree[node]) {       if (vis[x] == 0) {         vis[x] = 1;         dfs(vis, x, color);         blue += nodecolor[x][0];         red += nodecolor[x][1];         c += nodecolor[x][2];       }     }     nodecolor[node] = {blue, red, c};   }   int dfs2(vector<int>& vis, int node, int b, int r, int c, int ans) {     int b1 = 0, r1 = 0, b2 = 0, r2 = 0;     for (auto x : tree[node]) {       if (vis[x] == 0) {         vis[x] = 1;         b1 = b - (nodecolor[x][0]);         b2 = nodecolor[x][0];         r1 = r - (nodecolor[x][1]);         r2 = nodecolor[x][1];         if ((b1 == 0 || r1 == 0) && (b2 == 0 || r2 == 0)) {           ans += 1;         }         ans += dfs2(vis, x, b, r, c, 0);       }     }     return ans;   }   int main(int argc, const char* argv[]) {     ios_base::sync_with_stdio(false);     cin.tie(NULL);     int n, u, v;     cin >> n;     vector<int> color(n);     int b = 0, r = 0, c = 0;     for (int i = 0; i < n; i++) {       cin >> color[i];       b += (color[i] == 2);       c += (color[i] == 0);       r += (color[i] == 1);     }     for (int i = 0; i < n - 1; i++) {       cin >> u >> v;       tree[u].push_back(v);       tree[v].push_back(u);     }     int ans = 0;     vector<int> vis(n + 1, 0);     vis[1] = 1;     dfs(vis, 1, color);     for (int i = 0; i < n + 1; i++) {       vis[i] = 0;     }     vis[1] = 1;     ans = dfs2(vis, 1, b, r, c, 0);     cout << ans << endl;     return 0;   }   