func_code_string: #include <bits/stdc++.h>   using namespace std;   long long gcd(long long a, long long b) {     if (b > a) {       return gcd(b, a);     }     if (b == 0) {       return a;     }     return gcd(b, a % b);   }   long long expo(long long a, long long b, long long mod) {     long long res = 1;     while (b > 0) {       if (b & 1) res = (res * a) % mod;       a = (a * a) % mod;       b = b >> 1;     }     return res;   }   vector<int> sieve(int n) {     int *arr = new int[n + 1]();     vector<int> vect;     for (int i = 2; i <= n; i++)       if (arr[i] == 0) {         vect.push_back(i);         for (int j = 2 * i; j <= n; j += i) arr[j] = 1;       }     return vect;   }   void extendgcd(long long a, long long b, long long *v) {     if (b == 0) {       v[0] = 1;       v[1] = 0;       v[2] = a;       return;     }     extendgcd(b, a % b, v);     long long x = v[1];     v[1] = v[0] - v[1] * (a / b);     v[0] = x;     return;   }   long long mminv(long long a, long long b) {     long long arr[3];     extendgcd(a, b, arr);     return arr[0];   }   long long mminvprime(long long a, long long b) { return expo(a, b - 2, b); }   bool revsort(long long a, long long b) { return a > b; }   void swap(int &x, int &y) {     int temp = x;     x = y;     y = temp;   }   long long combination(long long n, long long r, long long m, long long *fact) {     long long val1 = fact[n];     long long val2 = mminvprime(fact[r], m);     long long val3 = mminvprime(fact[n - r], m);     return ((val1 * val2) % m * val3) % m;   }   void dfs(int start, vector<int> *edges, bool *visited, int *blues, int *reds,            int *color, vector<pair<int, int>> &childs) {     if (color[start] == 1) blues[start]++;     if (color[start] == 2) reds[start]++;     visited[start] = true;     for (int i = 0; i < edges[start].size(); i++) {       if (visited[edges[start][i]]) continue;       dfs(edges[start][i], edges, visited, blues, reds, color, childs);       blues[start] += blues[edges[start][i]];       reds[start] += reds[edges[start][i]];       childs.push_back({start, edges[start][i]});     }   }   int main() {     ios_base::sync_with_stdio(false);     cin.tie(NULL);     cout.tie(NULL);     int n;     cin >> n;     int color[n];     int sum1 = 0;     int sum2 = 0;     for (int i = 0; i < n; i++) {       cin >> color[i];       if (color[i] == 1) sum1++;       if (color[i] == 2) sum2++;     }     vector<int> *edges = new vector<int>[n];     vector<int> *ans1 = new vector<int>[n];     for (int i = 0; i < n - 1; i++) {       int u, v;       cin >> u >> v;       u--;       v--;       edges[u].push_back(v);       edges[v].push_back(u);       ans1[u].push_back(1);       ans1[v].push_back(1);     }     int *blues = new int[n]();     int *reds = new int[n]();     bool *visited = new bool[n]();     vector<pair<int, int>> childs;     dfs(0, edges, visited, blues, reds, color, childs);     int ans = 0;     set<pair<int, int>> s;     for (int i = 0; i < childs.size(); i++) {       int x = childs[i].second;       int blue = blues[x];       int red = reds[x];       if (blue == 0 && red == sum2) ans++;       if (red == 0 && blue == sum1) ans++;     }     cout << ans << endl;     return 0;   }   