func_code_string: #include <bits/stdc++.h>   using namespace std;   const int M = 998244353;   int n;   vector<int> v;   vector<set<int> > g;   int all[300011] = {};   set<pair<int, int> > ed;   pair<int, int> DFS(int u, int p) {     pair<int, int> ret(v[u], 1);     if (v[u] == 0) ret = {-1, 0};     for (auto v : g[u])       if (v != p) {         auto cur = DFS(v, u);         if (cur.first == -1) {           ed.insert({min(u, v), max(u, v)});           continue;         }         if (ret.first == -1 || cur.first == ret.first) {           ret.first = cur.first;           ret.second += cur.second;         } else {           puts( 0 );           exit(0);         }       }     if (ret.first == -1) return ret;     if (all[ret.first] == ret.second) ret = {-1, 0};     return ret;   }   vector<vector<int> > gr;   int w[300003];   int p[300003];   int find(int u) {     if (p[u] == u) return u;     return p[u] = find(p[u]);   }   int dp[300003][2];   int calc(int u, int p, int st) {     int &ret = dp[u][st];     if (ret != -1) return ret;     if (st == 0 && w[u] > 0) return ret = calc(u, p, 1);     ret = 0;     if (st == 1) {       ret = 1;       for (auto &v : gr[u])         if (v != p)           ret =               ret * (long long)(calc(v, u, 0) + calc(v, u, 1) * (w[v] == 0)) % M;     } else {       vector<int> pre(gr[u].size());       long long mul = 1;       for (int i = 0; i < gr[u].size(); ++i) {         int v = gr[u][i];         pre[i] = mul;         if (v == p) continue;         mul = mul * (calc(v, u, 0) + calc(v, u, 1) * (w[v] == 0)) % M;       }       mul = 1;       for (int i = (int)gr[u].size() - 1; i >= 0; --i) {         int v = gr[u][i];         if (v == p) continue;         ret = (ret + calc(v, u, 0) * mul % M * pre[i]) % M;         mul = mul * (calc(v, u, 0) + calc(v, u, 1) * (w[v] == 0)) % M;       }     }     return ret;   }   int main() {     int k;     scanf( %d%d , &n, &k);     for (int i = 0; i < n; ++i) p[i] = i;     v.resize(n);     g.resize(n);     for (auto &x : v) {       scanf( %d , &x);       ++all[x];     }     for (int i = 0; i < n; ++i) w[i] = v[i];     for (int i = 1, a, b; i < n; ++i) {       scanf( %d%d , &a, &b);       --a;       --b;       g[a].insert(b);       g[b].insert(a);     }     {       queue<int> q;       for (int i = 0; i < n; ++i)         if (v[i] == 0 && g[i].size() == 1) q.push(i);       while (!q.empty()) {         int u = q.front();         q.pop();         int v = *g[u].begin();         g[u].erase(v);         g[v].erase(u);         if (g[v].size() == 1 && w[v] == 0) q.push(v);       }     }     int root;     for (int i = 0; i < n; ++i)       if (g[i].size()) {         root = i;         break;       }     DFS(root, -1);     gr.resize(n);     for (int i = 0; i < g.size(); ++i)       for (auto x : g[i]) {         int u = i;         int v = x;         if (u > v) swap(u, v);         if (ed.find({u, v}) == ed.end()) {           u = find(u);           v = find(v);           p[u] = v;           w[v] = max(w[v], w[u]);         }       }     for (int i = 0; i < g.size(); ++i)       for (auto x : g[i]) {         int u = i;         int v = x;         if (u > v) continue;         if (ed.find({u, v}) != ed.end()) {           u = find(u);           v = find(v);           gr[u].push_back(v);           gr[v].push_back(u);         }       }     for (int i = 0; i < gr.size(); ++i)       if (gr[i].size() && w[i]) {         root = i;         break;       }     memset(dp, -1, sizeof(dp));     printf( %d n , calc(root, -1, 1));     return 0;   }   