#include <bits/stdc++.h>
using namespace std;
int n, k, h[300010], c[300010], f[300010], F[300010];
long long a[2][300010];
struct cmp {
  bool operator()(int x, int y) { return h[x] > h[y]; }
};
multiset<int, cmp> s[300010];
vector<int> p[300010], P[300010];
void dfs(int x) {
  for (auto i : p[x])
    if (i != f[x]) {
      f[i] = x;
      h[i] = h[x] + 1;
      dfs(i);
    }
}
void dp(int x) {
  a[0][x] = 1;
  for (auto i : P[x])
    if (F[x] != i) {
      F[i] = x;
      dp(i);
      a[1][x] = (a[1][x] * (a[0][i] + a[1][i]) + a[0][x] * a[1][i]) % 998244353;
      a[0][x] = a[0][x] * (a[0][i] + a[1][i]) % 998244353;
    }
  if (x <= k) a[1][x] = a[0][x], a[0][x] = 0;
}
int main() {
  scanf("%d%d", &n, &k);
  for (int i = (1); i <= (n); ++i) {
    scanf("%d", &c[i]);
  }
  for (int i = (1); i <= (n - 1); ++i) {
    int x, y;
    scanf("%d%d", &x, &y);
    p[x].push_back(y);
    p[y].push_back(x);
  }
  dfs(1);
  for (int i = (1); i <= (n); ++i)
    if (c[i]) s[c[i]].insert(i);
  for (int i = (1); i <= (k); ++i)
    while (s[i].size() > 1) {
      int x = *s[i].begin();
      s[i].erase(s[i].begin());
      if (!c[f[x]]) {
        c[f[x]] = i;
        s[i].insert(f[x]);
      }
      if (c[f[x]] != i) return printf("0\n"), 0;
    }
  int tot = k;
  for (int i = (1); i <= (n); ++i)
    if (!c[i]) c[i] = ++tot;
  for (int i = (1); i <= (n); ++i)
    for (auto j : p[i])
      if (c[i] != c[j]) {
        P[c[i]].push_back(c[j]);
      }
  dp(1);
  printf("%lld\n", a[1][1]);
}
