#include <bits/stdc++.h>
using namespace std;
const int maxn = 3e5 + 10;
const long long m = 998244353;
long long dp[maxn][2], r[maxn], l;
int col[maxn],
    flx[maxn] =
        {
            0,
},
    cnt[maxn] =
        {
            0,
},
    n, k;
int pre[maxn], a, b;
bool vis[maxn] = {
    0,
};
vector<int> v[maxn], w[maxn];
bool bfs() {
  int Q[maxn], front = 1, tail = 1;
  Q[front] = 1;
  vis[1] = true;
  int x, y, siz;
  while (front <= tail) {
    x = Q[front++];
    siz = v[x].size();
    for (int i = 0; i < siz; ++i) {
      y = v[x][i];
      if (!vis[y]) {
        Q[++tail] = y;
        vis[y] = true;
        pre[y] = x;
      }
    }
  }
  for (int i = n; i >= 1; --i) {
    x = Q[i];
    if (!col[x]) continue;
    if (cnt[col[x]] == 1) {
      flx[col[x]] = x;
      continue;
    }
    if (!col[pre[x]])
      col[pre[x]] = col[x];
    else if (col[pre[x]] == col[x])
      cnt[col[x]]--;
    else
      return false;
  }
  for (x = 2; x <= n; ++x) {
    y = pre[x];
    if (!col[x] || col[x] != col[y]) {
      if (col[y])
        w[flx[col[y]]].push_back(x);
      else
        w[y].push_back(x);
    }
  }
  return true;
}
void dfs(int x) {
  int siz = w[x].size(), y;
  for (int i = 0; i < siz; ++i) dfs(w[x][i]);
  r[siz] = 1;
  for (int i = siz - 1; i >= 0; --i) {
    y = w[x][i];
    r[i] = (r[i + 1] * (dp[y][0] + dp[y][1])) % m;
  }
  if (col[x]) {
    dp[x][0] = 0;
    dp[x][1] = r[0];
  } else {
    dp[x][0] = r[0];
    dp[x][1] = 0;
    l = 1;
    for (int i = 0; i < siz; ++i) {
      y = w[x][i];
      dp[x][1] = (dp[x][1] + (l * r[i + 1] % m) * dp[y][1]) % m;
      l = (l * (dp[y][0] + dp[y][1])) % m;
    }
  }
}
int main() {
  scanf("%d %d", &n, &k);
  for (int i = 1; i <= n; ++i) {
    scanf("%d", &col[i]);
    cnt[col[i]]++;
  }
  for (int i = 1; i < n; ++i) {
    scanf("%d %d", &a, &b);
    v[a].push_back(b);
    v[b].push_back(a);
  }
  if (!bfs()) {
    printf("0\n");
    return 0;
  }
  dfs(1);
  printf("%lld\n", dp[1][1]);
}
