#include <bits/stdc++.h>
using namespace std;
const int M = 3e5 + 5;
int n, k, a[M], head[M], cnt, has[M], siz[M], flag;
long long dp[M][2], suf[M];
struct edge {
  int v, next;
} e[M << 1];
void add(int u, int v) {
  e[++cnt].v = v;
  e[cnt].next = head[u];
  head[u] = cnt;
}
vector<int> v;
void dfs(int x, int fa) {
  siz[x] = (a[x] != 0);
  for (int i = head[x]; i != 0; i = e[i].next) {
    int to = e[i].v;
    if (to == fa) continue;
    dfs(to, x);
    if (flag) return;
    if (!a[x])
      a[x] = a[to];
    else if (a[to] && a[x] != a[to]) {
      flag = 1;
      return;
    }
    siz[x] += siz[to];
  }
  v.clear();
  for (int i = head[x]; i != 0; i = e[i].next) {
    int to = e[i].v;
    if (to == fa) continue;
    v.push_back(to);
  }
  int tot = v.size();
  if (!tot) {
    if (!a[x])
      dp[x][0] = 1;
    else
      dp[x][1] = 1;
    if (siz[x] == has[a[x]]) a[x] = siz[x] = 0;
    return;
  }
  suf[tot - 1] = 1;
  for (int to, i = tot - 2; i >= 0; i--)
    to = v[i + 1], suf[i] = suf[i + 1] * (dp[to][0] + dp[to][1]) % 998244353;
  if (a[x])
    dp[x][1] = suf[0] * (dp[v[0]][0] + dp[v[0]][1]) % 998244353;
  else {
    dp[x][0] = suf[0] * (dp[v[0]][0] + dp[v[0]][1]) % 998244353;
    long long pre = 1;
    for (int i = 0; i <= tot - 1; i++) {
      int to = v[i];
      dp[x][1] = (dp[x][1] + pre * dp[to][1] % 998244353 * suf[i] % 998244353) %
                 998244353;
      pre = pre * (dp[to][0] + dp[to][1]) % 998244353;
    }
  }
  if (siz[x] == has[a[x]]) a[x] = siz[x] = 0;
}
int main() {
  scanf("%d%d", &n, &k);
  for (int i = 1; i <= n; i++) scanf("%d", &a[i]), has[a[i]]++;
  for (int x, y, i = 1; i <= n - 1; i++) {
    scanf("%d%d", &x, &y);
    add(x, y);
    add(y, x);
  }
  dfs(1, 0);
  if (flag)
    printf("0\n");
  else
    printf("%lld\n", dp[1][1] % 998244353);
  return 0;
}
