func_code_string: #include <bits/stdc++.h>   using namespace std;   using ll = long long;   const int MXN = 3e5 + 5;   const int MOD = 998244353;   const int LOG = 19;   int add(int x, int y) { return (x += y) < MOD ? x : x - MOD; }   int sub(int x, int y) { return (x -= y) >= 00 ? x : x + MOD; }   int mul(int x, int y) { return (ll)x * y % MOD; }   int pwr(int x, int p) {     if (p == 0) return 1;     if (p % 2) return mul(x, pwr(x, p - 1));     return pwr(mul(x, x), p / 2);   }   int inv(int x) { return pwr(x, MOD - 2); }   int sol = 1;   int ok[MXN];   int val[MXN];   int dep[MXN];   int lca[MXN];   int nxt[MXN];   int dp[MXN][2];   int anc[MXN][LOG];   vector<int> adj[MXN];   int getlca(int x, int y) {     if (dep[x] < dep[y]) swap(x, y);     int diff = dep[x] - dep[y];     for (int i = 0; i < LOG; ++i) {       if (diff >> i & 1) x = anc[x][i];     }     if (x == y) return x;     for (int i = LOG - 1; i >= 0; --i) {       if (anc[x][i] != anc[y][i]) {         x = anc[x][i];         y = anc[y][i];       }     }     return anc[x][0];   }   void dfs1(int x, int p) {     if (val[x]) {       if (lca[val[x]]) {         lca[val[x]] = getlca(lca[val[x]], x);       } else {         lca[val[x]] = x;       }     }     for (int y : adj[x]) {       if (y == p) continue;       anc[y][0] = x;       dep[y] = dep[x] + 1;       for (int i = 0; i < LOG - 1; ++i) {         anc[y][i + 1] = anc[anc[y][i]][i];       }       dfs1(y, x);     }   }   void dfs2(int x, int p) {     vector<int> vec;     for (int y : adj[x]) {       if (y == p) continue;       dfs2(y, x);       if (ok[y]) {         ok[x] = 1;         vec.push_back(y);       }       if (val[y] && y != lca[val[y]]) {         if (val[x] && val[x] != val[y]) {           cout <<  0 n ;           exit(0);         }         val[x] = val[y];       }     }     if (val[x]) {       dp[x][0] = ok[x] = 1;       for (int y : vec) {         sol = mul(sol, add(mul(dp[y][0], nxt[y] + 1), dp[y][1]));       }     } else if (vec.size() > 1) {       dp[x][1] = 1;       for (int y : vec) {         int tmp = add(dp[y][1], mul(dp[y][0], nxt[y] + 1));         dp[x][1] = mul(dp[x][1], tmp);       }       for (int y : vec) {         int tmp = add(dp[y][1], mul(dp[y][0], nxt[y] + 1));         dp[x][0] = add(dp[x][0], mul(mul(dp[x][1], inv(tmp)), dp[y][0]));       }     } else if (vec.size() == 1) {       dp[x][0] = dp[vec[0]][0];       dp[x][1] = dp[vec[0]][1];       nxt[x] = nxt[vec[0]] + 1;     }   }   int main() {     ios::sync_with_stdio(0), cin.tie(0);     int n, k;     cin >> n >> k;     for (int i = 1; i <= n; ++i) {       cin >> val[i];     }     for (int i = 2; i <= n; ++i) {       int x, y;       cin >> x >> y;       adj[x].push_back(y);       adj[y].push_back(x);     }     dfs1(1, 0);     dfs2(1, 0);     if (dp[1][0]) sol = mul(sol, dp[1][0]);     cout << sol <<   n ;   }   