#include <bits/stdc++.h>
using namespace std;
const int C = 300001, M = 998244353;
vector<vector<int> > tr(C);
int ij[C];
const int _LCA_D = 21;
int s[C], is = 0, ip = 1, pre[C], apre[C], dep[C], son[C], par[_LCA_D][C],
          ji[C];
void proc_tree(int a, int n, vector<vector<int> > &tab, int ij[]) {
  s[0] = a, is = 1, pre[a] = 1, apre[a] = 1, dep[a] = 1, ip = 2;
  int b;
  for (int z = 1; z <= n; z++) ji[z] = ij[z];
  while (is > 0) {
    a = s[is - 1];
    if (ji[a] > 0 && tab[a][ji[a] - 1] == par[0][a]) ji[a]--;
    if (ji[a] > 0)
      b = s[is] = tab[a][ji[a] - 1], pre[b] = ip, apre[ip] = b, par[0][b] = a,
      dep[b] = is + 1, ip++, is++, ji[a]--;
    else
      son[a]++, son[par[0][a]] += son[a], is--;
  }
}
void prep_lca(int n, int f = _LCA_D) {
  par[0][1] = 1;
  for (int j = 1; j < f; j++) {
    for (int i = 1; i <= n; i++) {
      par[j][i] = par[j - 1][par[j - 1][i]];
    }
  }
}
int anc(int a, int b) {
  if (pre[a] <= pre[b] && pre[a] + son[a] > pre[b]) return 1;
  return 0;
}
int lca(int a, int b, int f = _LCA_D) {
  if (anc(a, b) == 1) return a;
  if (anc(b, a) == 1) return b;
  int i = f - 1, x = a;
  while (i >= 0) {
    if (anc(par[i][x], b) == 1)
      i--;
    else
      x = par[i][x];
  }
  return par[0][x];
}
int color[C], lowest[C];
long long dp[C][2];
int main() {
  int n, k, res, i, a, b, pa;
  scanf("%d %d", &n, &k);
  for (i = 1; i <= n; i++) scanf("%d", &color[i]);
  for (int z = 1; z < n; z++)
    scanf("%d %d", &a, &b), tr[a].push_back(b), tr[b].push_back(a), ij[a]++,
        ij[b]++;
  proc_tree(1, n, tr, ij);
  prep_lca(n);
  for (i = 1; i <= k; i++) lowest[i] = -1;
  for (i = 1; i <= n; i++) {
    if (color[i] == 0) continue;
    if (lowest[color[i]] == -1)
      lowest[color[i]] = i;
    else
      lowest[color[i]] = lca(i, lowest[color[i]]);
  }
  for (i = n; i > 0; i--) {
    a = apre[i];
    if (color[a] == 0 || a == 1) continue;
    if (lowest[color[a]] != a) {
      pa = par[0][a];
      if (color[pa] != color[a] && color[pa] != 0) {
        printf("0\n");
        return 0;
      }
      color[pa] = color[a];
    }
  }
  for (i = 1; i <= n; i++) dp[i][0] = 1, dp[i][1] = 0;
  for (i = n; i > 0; i--) {
    a = apre[i];
    pa = par[0][a];
    if (color[a] != 0) dp[a][1] = 0;
    if (a == lowest[color[a]]) dp[a][1] = dp[a][0], dp[a][0] = 0;
    if (i == 1) break;
    dp[pa][1] = (dp[pa][1] * (dp[a][0] + dp[a][1]) + dp[pa][0] * dp[a][1]) % M;
    dp[pa][0] = (dp[pa][0] * (dp[a][0] + dp[a][1])) % M;
  }
  printf("%lld\n", dp[1][1]);
  return 0;
}
