#include <bits/stdc++.h>
using namespace std;
int n, k;
int c[300010], lc[300010], now[300010];
int fa[30][300010], dep[300010], g[300010];
long long f[300010][2];
bool flag;
vector<int> E[300010], Q1[300010], Q2[300010];
vector<long long> pre, suf;
inline void insert(int a, int b) {
  E[a].push_back(b);
  E[b].push_back(a);
}
int lca(int a, int b) {
  int k;
  if (dep[a] != dep[b]) {
    if (dep[a] < dep[b]) swap(a, b);
    k = 0;
    while (dep[fa[k][a]] > dep[b]) k++;
    while (k >= 0) {
      if (dep[fa[k][a]] > dep[b]) a = fa[k][a];
      k--;
    }
    a = fa[0][a];
  }
  if (a == b) return a;
  k = 0;
  while (fa[k][a] != fa[k][b]) k++;
  while (k >= 0) {
    if (fa[k][a] != fa[k][b]) {
      a = fa[k][a];
      b = fa[k][b];
    }
    k--;
  }
  return fa[0][a];
}
void dfs1(int x) {
  for (int i = 0; i < E[x].size(); i++)
    if (fa[0][x] != E[x][i]) {
      fa[0][E[x][i]] = x;
      dep[E[x][i]] = dep[x] + 1;
      dfs1(E[x][i]);
    }
}
void dfs2(int x) {
  if (c[x]) {
    if (now[c[x]])
      g[x] = c[x];
    else
      g[x] = -1;
  }
  for (int i = 0; i < E[x].size(); i++)
    if (fa[0][x] != E[x][i]) {
      now[c[x]]++;
      dfs2(E[x][i]);
      if (flag) return;
      now[c[x]]--;
      if (g[E[x][i]] && !g[x]) g[x] = -1;
      if (g[E[x][i]] > 0 && now[g[E[x][i]]]) {
        if (c[x] && g[E[x][i]] != c[x]) {
          flag = 1;
          return;
        }
        if (!g[x] || g[x] == -1) {
          g[x] = g[E[x][i]];
          c[x] = g[x];
          continue;
        }
        if (g[x] != g[E[x][i]]) {
          flag = 1;
          return;
        }
      }
    }
}
void dfs3(int x) {
  int a;
  long long s;
  for (int i = 0; i < E[x].size(); i++)
    if (fa[0][x] != E[x][i]) {
      dfs3(E[x][i]);
      if (g[E[x][i]] > 0) Q1[x].push_back(E[x][i]);
      if (g[E[x][i]] == -1) Q2[x].push_back(E[x][i]);
    }
  if (Q1[x].empty() && Q2[x].empty()) {
    if (c[x]) {
      f[x][0] = 0;
      f[x][1] = 1;
    } else {
      f[x][0] = 1;
      f[x][1] = 0;
    }
    return;
  }
  if (!c[x] && Q1[x].empty()) {
    f[x][0] = s = 1;
    a = Q2[x].size();
    for (int i = 0; i < a; i++) {
      f[x][0] = f[x][0] * (f[Q2[x][i]][0] + f[Q2[x][i]][1]) % 998244353;
      pre.push_back(f[x][0]);
    }
    for (int i = a - 1; i >= 0; i--) {
      s = s * (f[Q2[x][i]][0] + f[Q2[x][i]][1]) % 998244353;
      suf.push_back(s);
    }
    reverse(suf.begin(), suf.end());
    for (int i = 0; i < a; i++) {
      s = f[Q2[x][i]][1];
      if (i) s = s * pre[i - 1] % 998244353;
      if (i < a - 1) s = s * suf[i + 1] % 998244353;
      f[x][1] += s;
      if (f[x][1] >= 998244353) f[x][1] %= 998244353;
    }
    pre.clear();
    suf.clear();
  } else {
    f[x][0] = 0;
    f[x][1] = 1;
    for (int i = 0; i < Q2[x].size(); i++)
      f[x][1] = f[x][1] * (f[Q2[x][i]][0] + f[Q2[x][i]][1]) % 998244353;
    for (int i = 0; i < Q1[x].size(); i++)
      f[x][1] = f[x][1] * f[Q1[x][i]][1] % 998244353;
  }
}
int main() {
  int a, b;
  scanf("%d%d", &n, &k);
  for (int i = 1; i <= n; i++) scanf("%d", &c[i]);
  for (int i = 1; i < n; i++) {
    scanf("%d%d", &a, &b);
    insert(a, b);
  }
  dfs1(1);
  for (int i = 1; i <= 20; i++)
    for (int j = 1; j <= n; j++) fa[i][j] = fa[i - 1][fa[i - 1][j]];
  for (int i = 1; i <= n; i++)
    if (c[i]) {
      if (!lc[c[i]]) {
        lc[c[i]] = i;
        continue;
      }
      lc[c[i]] = lca(lc[c[i]], i);
      if (!c[lc[c[i]]]) {
        c[lc[c[i]]] = c[i];
        continue;
      }
      if (c[lc[c[i]]] != c[i]) {
        puts("0");
        return 0;
      }
    }
  dfs2(1);
  if (flag) {
    puts("0");
    return 0;
  }
  dfs3(1);
  cout << f[1][1] << endl;
  return 0;
}
