#include <bits/stdc++.h>
using namespace std;
template <typename T>
void maxtt(T& t1, T t2) {
  t1 = max(t1, t2);
}
template <typename T>
void mintt(T& t1, T t2) {
  t1 = min(t1, t2);
}
bool debug = 0;
int n, m, k;
int dx[4] = {0, 1, 0, -1}, dy[4] = {1, 0, -1, 0};
string direc = "URDL";
long long ln, lk, lm;
void etp(bool f = 0) {
  puts(f ? "YES" : "NO");
  exit(0);
}
void addmod(int& x, int y, int mod = 998244353) {
  assert(y >= 0);
  x += y;
  if (x >= mod) x -= mod;
  assert(x >= 0 && x < mod);
}
void et(int x = -1) {
  printf("%d\n", x);
  exit(0);
}
long long fastPow(long long x, long long y, int mod = 998244353) {
  long long ans = 1;
  while (y > 0) {
    if (y & 1) ans = (x * ans) % mod;
    x = x * x % mod;
    y >>= 1;
  }
  return ans;
}
long long gcd1(long long x, long long y) { return y ? gcd1(y, x % y) : x; }
vector<int> mp[300135];
int a[300135];
void ppt(int ans = 0) {
  printf("%d\n", ans);
  exit(0);
}
int d[300135], p[300135][20];
vector<int> pos[300135];
void dfs(int x, int pa) {
  d[x] = d[pa] + 1;
  p[x][0] = pa;
  for (int v : mp[x])
    if (v != pa) {
      dfs(v, x);
    }
}
void init() {
  for (int j = 1; (1 << j) <= n; j++)
    for (int i = 1; i <= n; i++) p[i][j] = p[p[i][j - 1]][j - 1];
}
int lca(int a, int b) {
  if (d[a] > d[b]) swap(a, b);
  int f = d[b] - d[a];
  for (int i = 0; (1 << i) <= f; i++)
    if ((1 << i) & f) b = p[b][i];
  if (a != b) {
    for (int i = 19; i >= 0; i--)
      if (p[a][i] != p[b][i]) {
        a = p[a][i];
        b = p[b][i];
      }
    a = p[a][0];
  }
  return a;
}
int hd[300135], dp[300135][2];
int cal(int x, int pa) {
  vector<pair<int, int> > cs;
  int c1 = 0;
  for (int c : mp[x])
    if (c != pa) {
      int z = cal(c, x);
      if (z != 0) cs.push_back({z, c});
      if (z == -1) c1++;
    }
  if (cs.empty()) {
    if (a[x] == 0) {
      dp[x][0] = 1;
    } else
      dp[x][1] = 1;
    if (a[x] == 0) return 0;
    return x == hd[a[x]] ? -1 : a[x];
  }
  int N = cs.size();
  vector<int> ss(N + 5, 0), pre(N + 5, 0), nxt(N + 5, 0);
  pre[0] = 1;
  nxt[N + 1] = 1;
  for (int(i) = 0; (i) < (int)(N); (i)++) {
    addmod(ss[i + 1], dp[cs[i].second][0]);
    addmod(ss[i + 1], dp[cs[i].second][1]);
    pre[i + 1] = (long long)pre[i] * ss[i + 1] % 998244353;
  }
  for (int i = N; i; i--) {
    nxt[i] = (long long)nxt[i + 1] * ss[i] % 998244353;
  }
  if (c1 == cs.size() && a[x] == 0) {
    dp[x][0] = pre[N];
    for (int(i) = 1; (i) <= (int)(N); (i)++) {
      addmod(dp[x][1], (long long)dp[cs[i - 1].second][1] * pre[i - 1] %
                           998244353 * nxt[i + 1] % 998244353);
    }
  } else {
    assert(a[x] != 0);
    int A = 1, B = 1;
    for (int(i) = 1; (i) <= (int)(N); (i)++) {
      if (cs[i - 1].first == -1)
        A = (long long)A * ss[i] % 998244353;
      else {
        assert(cs[i - 1].first == a[x]);
        B = (long long)B * dp[cs[i - 1].second][1] % 998244353;
      }
    }
    dp[x][1] = (long long)A * B % 998244353;
  }
  return a[x] == 0 || hd[a[x]] == x ? -1 : a[x];
}
void fmain(int tid) {
  scanf("%d%d", &n, &k);
  for (int(i) = 1; (i) <= (int)(n); (i)++) {
    scanf("%d", a + i);
    pos[a[i]].push_back(i);
  }
  for (int i = 0, u, v; i < n - 1; i++) {
    scanf("%d%d", &u, &v);
    mp[u].push_back(v);
    mp[v].push_back(u);
  };
  dfs(1, 0);
  init();
  for (int(i) = 1; (i) <= (int)(k); (i)++) {
    if (pos[i].size() == 1) {
      hd[i] = pos[i][0];
      continue;
    }
    int f = lca(pos[i][0], pos[i][1]);
    for (int j = 2; j < pos[i].size(); j++) f = lca(f, pos[i][j]);
    if (a[f] != 0 && a[f] != i) ppt();
    a[f] = i;
    hd[i] = f;
    for (int z : pos[i])
      if (d[z] > d[f]) {
        int x = p[z][0];
        while (a[x] == 0) {
          a[x] = i;
          x = p[x][0];
        }
        if (a[x] != i) ppt();
      }
  }
  cal(1, 0);
  printf("%d\n", dp[1][1]);
}
int main() {
  int t = 1;
  for (int(i) = 1; (i) <= (int)(t); (i)++) {
    fmain(i);
  }
  return 0;
}
