#include <bits/stdc++.h>
using namespace std;
const int N = 3e5 + 5;
int a[N];
vector<int> e[N];
bool vis[N];
int col[N];
class Lca {
  int f[32][N << 1], lg[N << 1], p[32];
  int tot;
  int mi(int a, int b) { return d[a] < d[b] ? a : b; }
  void dfs(int u, int dep, int fa) {
    ver[++tot] = u;
    first[u] = tot;
    d[tot] = dep;
    pre[u] = fa;
    for (auto v : e[u])
      if (v != fa) {
        dfs(v, dep + 1, u);
        ver[++tot] = u, d[tot] = dep;
      }
  }
  int query(int l, int r) {
    r++;
    int t = lg[r - l];
    r -= p[t];
    return ver[mi(f[t][l], f[t][r])];
  }

 public:
  int d[N << 1], ver[N << 1], first[N << 1], pre[N << 1];
  void init() {
    dfs(1, tot = 0, 0), p[0] = 1, lg[0] = -1;
    int n = tot;
    for (int i = 1; (1 << i) <= n; ++i) p[i] = p[i - 1] * 2;
    for (int i = 1; i <= n; ++i) lg[i] = lg[i / 2] + 1, f[0][i] = i;
    for (int i = 1; (1 << i) <= n; ++i)
      for (int j = 1; j + p[i] - 1 <= n; ++j) {
        f[i][j] = mi(f[i - 1][j], f[i - 1][j + p[i - 1]]);
      }
  }
  int lca(int u, int v) {
    int x = first[u], y = first[v];
    if (x > y) swap(x, y);
    return query(x, y);
  }
} Q;
vector<int> c[N], g2[N];
int f[N];
int find(int u) { return u == f[u] ? u : f[u] = find(f[u]); }
void knit(int v, int u) {
  v = find(v), u = find(u);
  f[v] = u;
}
const int mod = 998244353;
long long dp[N][2];
int id[N];
int n, k;
void dfs(int u, int fa) {
  dp[u][u <= k] = 1;
  vis[u] = 1;
  for (auto v : g2[u]) {
    if (!vis[v]) {
      dfs(v, u);
      dp[u][1] = (dp[u][1] * dp[v][0] + dp[u][0] * dp[v][1] +
                  dp[u][1] * dp[v][1] % mod) %
                 mod;
      dp[u][0] = (dp[u][0] * dp[v][0] + dp[u][0] * dp[v][1]) % mod;
    }
  }
}
int main() {
  scanf("%d%d", &n, &k);
  for (int i = 1; i <= n; ++i) scanf("%d", &a[i]), c[a[i]].push_back(i);
  for (int i = 1; i < n; ++i) {
    int u, v;
    scanf("%d%d", &u, &v);
    e[u].push_back(v), e[v].push_back(u);
  }
  Q.init();
  int tot = 0, fail = 0;
  for (int i = 1; i <= n; ++i) f[i] = i;
  for (int i = 1; i <= k; ++i) {
    tot = fail = 0;
    if (!c[i].size()) continue;
    sort(c[i].begin(), c[i].end(),
         [](int a, int b) { return Q.first[a] < Q.first[b]; });
    int tmp = c[i][0];
    for (int j = 1; j < c[i].size(); ++j) {
      int lca = Q.lca(c[i][j], c[i][j - 1]);
      if (j == 1 || Q.d[Q.first[lca]] < Q.d[Q.first[tmp]]) tmp = lca;
    }
    for (auto u : c[i]) {
      for (int v = u; v != tmp; v = f[v]) {
        if (a[v] != 0 && a[v] != i) {
          fail = 1;
          break;
        }
        col[v] = i;
        knit(v, Q.pre[v]);
      }
    }
    if (a[tmp] != 0 && a[tmp] != i) fail = 1;
    col[tmp] = i;
    if (fail) break;
  }
  if (fail)
    puts("0");
  else {
    tot = 0;
    for (int i = 1; i <= n; ++i) id[i] = col[i] ? col[i] : (++tot) + k;
    for (int i = 1; i <= n; ++i) {
      for (auto v : e[i]) {
        int x = id[v], y = id[i];
        if (x != y) g2[x].push_back(y);
      }
    }
    dfs(1, 0);
    cout << dp[1][1] << endl;
  }
  return 0;
}
