func_code_string: #include <bits/stdc++.h>   using namespace std;   using ll = long long;   template <class T, class U>   ostream& operator<<(ostream& o, const pair<T, U>& p) {     o <<  (  << p.first <<  ,  << p.second <<  ) ;     return o;   }   template <class T>   ostream& operator<<(ostream& o, const vector<T>& v) {     o <<  [ ;     for (T t : v) {       o << t <<  , ;     }     o <<  ] ;     return o;   }   const int MAX_V = 300003;   const int MAX_LOG_V = 19;   vector<int> G[MAX_V];   int root = 0;   int parent[MAX_LOG_V][MAX_V];   int depth[MAX_V];   void lca_dfs(int v, int p, int d) {     parent[0][v] = p;     depth[v] = d;     for (int i = 0; i < (int)(G[v].size()); ++i) {       if (G[v][i] != p) {         lca_dfs(G[v][i], v, d + 1);       }     }   }   void lca_init(int V) {     lca_dfs(root, -1, 0);     for (int k = 0; k + 1 < MAX_LOG_V; ++k) {       for (int v = 0; v < V; ++v) {         if (parent[k][v] < 0)           parent[k + 1][v] = -1;         else           parent[k + 1][v] = parent[k][parent[k][v]];       }     }   }   int lca(int u, int v) {     if (depth[u] > depth[v]) swap(u, v);     for (int k = 0; k < MAX_LOG_V; ++k) {       if ((depth[v] - depth[u]) >> k & 1) v = parent[k][v];     }     if (u == v) return u;     for (int k = MAX_LOG_V - 1; k >= 0; --k) {       if (parent[k][u] != parent[k][v]) {         u = parent[k][u];         v = parent[k][v];       }     }     return parent[0][u];   }   const int N = MAX_V;   const ll mod = 998244353;   int n, k;   int a[N];   int L[N];   bool invalid = false;   ll dp[N][2] = {};   int dfs(int x, int par) {     map<int, int> m;     for (int e : G[x])       if (e != par) {         int res = dfs(e, x);         if (res != 0) m[e] = res;       }     if (m.empty()) {       int idx = (a[x] != 0);       dp[x][idx] = 1;       if (L[a[x]] == x) return -1;       return a[x];     }     unordered_set<int> s;     vector<int> ch;     for (const auto& p : m) {       ch.push_back(p.first);       if (p.second != -1) s.insert(p.second);     }     if (s.size() > 1) {       invalid = true;       return 0;     }     if (a[x] == 0 && s.size() == 0) {       dp[x][0] = 1;       for (int v : ch) (dp[x][0] *= (dp[v][0] + dp[v][1])) %= mod;       int C = ch.size();       vector<ll> pre(C), suf(C);       for (int i = 0; i < (int)(C); ++i)         pre[i] = suf[i] = (dp[ch[i]][0] + dp[ch[i]][1]) % mod;       for (int i = 0; i < (int)(C - 1); ++i) (pre[i + 1] *= pre[i]) %= mod;       for (int i = C - 1; i > 0; --i) (suf[i - 1] *= suf[i]) %= mod;       dp[x][1] = 0;       for (int i = 0; i < (int)(C); ++i) {         ll mul = 1;         if (i != 0) (mul *= pre[i - 1]) %= mod;         if (i != C - 1) (mul *= suf[i + 1]) %= mod;         ll add = (mul * dp[ch[i]][1]) % mod;         (dp[x][1] += add) %= mod;       }       return -1;     } else {       int col = a[x];       if (!s.empty()) {         assert(s.size() == 1);         int h = *s.begin();         if (col != 0 && col != h) {           invalid = true;           return 0;         }         col = h;       }       dp[x][1] = 1;       for (int v : ch) {         if (m[v] == -1)           (dp[x][1] *= (dp[v][0] + dp[v][1])) %= mod;         else           (dp[x][1] *= dp[v][1]) %= mod;       }       return (x == L[col]) ? -1 : col;     }   }   int main() {     scanf(  %d %d , &n, &k);     vector<vector<int>> col(k + 1);     for (int i = 0; i < (int)(n); ++i) {       scanf(  %d , &a[i]);       col[a[i]].push_back(i);     }     for (int i = 0; i < (int)(n - 1); ++i) {       int u, v;       scanf(  %d %d , &u, &v);       --u;       --v;       G[u].push_back(v);       G[v].push_back(u);     }     lca_init(n);     for (int i = 1; i <= k; ++i) {       int x = col[i][0];       for (int j : col[i]) x = lca(x, j);       if (a[x] != 0 && a[x] != i) invalid = true;       a[x] = i;       L[i] = x;     }     dfs(0, -1);     printf( %lld n , invalid ? 0 : dp[0][1]);     return 0;   }   