func_code_string: #include <bits/stdc++.h>   using namespace std;   int n, k, t[300010], u, v, a[300010], conf[300010], tc[300010], check[300010];   long long dp[300010][2];   vector<int> edge[300010];   const long long MOD = 998244353;   long long add(long long a, long long b) {     a += b;     if (a >= MOD) a -= MOD;     return a;   }   long long times(long long a, long long b) {     a *= b;     a %= MOD;     return a;   }   long long fexp(long long m) {     long long h = MOD - 2ll;     long long r = 1ll;     while (h > 0) {       m %= MOD;       if (h & 1ll) r *= m;       m *= m;       h >>= 1ll;       r %= MOD;     }     return r;   }   void conect(int vert, int dis) {     for (int i = 0; i < t[dis]; i++)       if (edge[dis][i] != vert) {         edge[vert].push_back(edge[dis][i]);         t[vert]++;         for (int j = 0; j < t[edge[dis][i]]; j++)           if (edge[edge[dis][i]][j] == dis) {             edge[edge[dis][i]][j] = vert;             break;           }       }     t[dis] = 0;   }   void super_node(int vert, int color) {     conf[vert] = -1;     check[vert] = a[vert];     int h;     if (a[vert] > 0) color = a[vert];     for (int i = 0; i < t[vert]; i++)       if (conf[edge[vert][i]] >= 0) {         if (conf[edge[vert][i]] == 0) super_node(edge[vert][i], color);         if (t[edge[vert][i]] == 1 && a[edge[vert][i]] == 0) {           t[edge[vert][i]] = 0;           swap(edge[vert][i], edge[vert][t[vert] - 1]);           edge[vert].pop_back();           t[vert]--;           i--;         } else if (check[edge[vert][i]] == color) {           tc[a[edge[vert][i]]]--;           conect(vert, edge[vert][i]);           swap(edge[vert][i], edge[vert][t[vert] - 1]);           edge[vert].pop_back();           t[vert]--;           i--;           if (a[vert] == 0) {             check[vert] = color;             a[vert] = color;             tc[color]++;           }         } else if (a[vert] == 0) {           if (check[vert] == check[edge[vert][i]]) {             color = check[vert];             tc[a[edge[vert][i]]]--;             conect(vert, edge[vert][i]);             swap(edge[vert][i], edge[vert][t[vert] - 1]);             edge[vert].pop_back();             i--;             t[vert]--;             swap(edge[vert][h], edge[vert][i]);             tc[a[edge[vert][i]]]--;             conect(vert, edge[vert][i]);             swap(edge[vert][i], edge[vert][t[vert] - 1]);             edge[vert].pop_back();             i--;             t[vert]--;             a[vert] = color;             tc[color]++;           } else if (check[vert] == 0) {             check[vert] = check[edge[vert][i]];             h = i;           } else if (tc[check[edge[vert][i]]] > 1 && check[vert] != color) {             check[vert] = check[edge[vert][i]];             h = i;           }         }       }     conf[vert] = 1;   }   void resp(int vert) {     conf[vert] = -1;     if (a[vert] == 0)       dp[vert][0] = 1;     else       dp[vert][1] = 1;     for (int i = 0; i < t[vert]; i++)       if (conf[edge[vert][i]] == 0) {         resp(edge[vert][i]);         dp[vert][0] =             times(dp[vert][0], add(dp[edge[vert][i]][0], dp[edge[vert][i]][1]));         dp[vert][1] =             times(dp[vert][1], add(dp[edge[vert][i]][0], dp[edge[vert][i]][1]));       }     if (a[vert] == 0) {       long long hf = 1;       for (int i = 0; i < t[vert]; i++)         if (conf[edge[vert][i]] == 1)           hf = times(hf, add(dp[edge[vert][i]][0], dp[edge[vert][i]][1]));       for (int i = 0; i < t[vert]; i++)         if (conf[edge[vert][i]] == 1) {           long long h = dp[edge[vert][i]][1];           h = times(h, hf);           h = times(h, fexp(add(dp[edge[vert][i]][0], dp[edge[vert][i]][1])));           dp[vert][1] = add(dp[vert][1], h);         }     }     conf[vert] = 1;   }   int main() {     scanf( %d %d , &n, &k);     for (int i = 1; i <= n; i++) {       scanf( %d , &a[i]);       tc[a[i]]++;     }     tc[0] = 0;     for (int i = 1; i < n; i++) {       scanf( %d %d , &u, &v);       edge[u].push_back(v);       edge[v].push_back(u);       t[u]++;       t[v]++;     }     for (int i = 1; i <= n; i++)       if (a[i] > 0) {         super_node(i, a[i]);         break;       }     for (int i = 1; i <= k; i++)       if (tc[i] != 1) {         printf( 0 n );         return 0;       }     for (int i = 1; i <= n; i++) conf[i] = 0;     for (int i = 1; i <= n; i++)       if (t[i] > 0) {         resp(i);         printf( %lld n , dp[i][1]);         return 0;       }   }   