#include <bits/stdc++.h>
using namespace std;
template <typename T>
inline T read(register T& t) {
  register T f = 1;
  register char ch = getchar();
  t = 0;
  while (ch < '0' || ch > '9') {
    if (ch == '-') f = -f;
    ch = getchar();
  }
  while (ch >= '0' && ch <= '9') t = t * 10 + ch - '0', ch = getchar();
  t *= f;
  return t;
}
template <typename T, typename... Args>
inline void read(T& t, Args&... args) {
  read(t);
  read(args...);
}
const long long p = 998244353;
inline long long power(register long long x, register long long k = p - 2) {
  register long long re = 1;
  for (; k; k >>= 1, x = x * x % p)
    if (k & 1) re = re * x % p;
  return re;
}
int n, k;
int col[300005];
int head[300005], to[600005], ne[600005], total = 1;
inline void add(int a, int b) {
  total++;
  to[total] = b;
  ne[total] = head[a];
  head[a] = total;
}
int deep[300005], size[300005];
int son[300005], fa[300005];
int top[300005], id[300005], id_top;
void dfs1(int now) {
  deep[now] = deep[fa[now]] + 1;
  size[now] = 1;
  for (int i = head[now]; i; i = ne[i])
    if (to[i] != fa[now]) {
      fa[to[i]] = now;
      dfs1(to[i]);
      size[now] += size[to[i]];
      if (size[to[i]] > size[son[now]]) son[now] = to[i];
    }
}
void dfs2(int now, int topf) {
  id[now] = ++id_top;
  top[now] = topf;
  if (son[now]) dfs2(son[now], topf);
  for (int i = head[now]; i; i = ne[i])
    if (to[i] != fa[now] && to[i] != son[now]) dfs2(to[i], to[i]);
}
inline int lca(int x, int y) {
  while (top[x] != top[y]) {
    if (deep[top[x]] < deep[top[y]]) swap(x, y);
    x = fa[top[x]];
  }
  if (deep[x] < deep[y]) return x;
  return y;
}
vector<int> pos[300005];
long long f[300005][2];
void dfs(int now, int fa) {
  long long mul = 1;
  for (int i = head[now]; i; i = ne[i])
    if (to[i] != fa) {
      dfs(to[i], now);
      mul = mul * (f[to[i]][0] + f[to[i]][1]) % p;
    }
  if (col[now] == 0) {
    f[now][0] = mul;
    for (int i = head[now]; i; i = ne[i])
      if (to[i] != fa)
        f[now][1] = (f[now][1] + mul * power((f[to[i]][0] + f[to[i]][1]) % p) %
                                     p * f[to[i]][1]) %
                    p;
  } else
    f[now][1] = mul;
}
int main() {
  read(n, k);
  for (int i = 1; i <= n; i++) read(col[i]), pos[col[i]].push_back(i);
  for (int i = 1, x, y; i < n; i++) read(x, y), add(x, y), add(y, x);
  dfs1(1);
  dfs2(1, 1);
  for (int i = 1; i <= k; i++) {
    if (pos[i].size() <= 1) continue;
    int lc = pos[i][0];
    for (int j : pos[i]) lc = lca(lc, j);
    if (col[lc] && col[lc] != i) return printf("0\n"), 0;
    col[lc] = i;
    for (int j : pos[i])
      while (1) {
        if (j == lc) break;
        j = fa[j];
        if (j == 0) break;
        if (col[j] == i) break;
        if (col[j]) return printf("0\n"), 0;
        col[j] = i;
      }
  }
  dfs(1, 1);
  printf("%lld\n", f[1][1]);
  return 0;
}
