#include <bits/stdc++.h>
using namespace std;
const int MXLOG = 30;
const long long mod = 998244353;
const int MX = 1e6 + 5;
int col[MX], de[MX], fa[MX], tp[MX], TP[MX];
vector<int> g[MX], ng[MX], same[MX];
long long dp[MX][2], pa[MX], suf[MX];
vector<int> have;
void dfs(int now, int last) {
  for (auto i : g[now]) {
    if (i == last) continue;
    de[i] = de[now] + 1, fa[i] = now;
    dfs(i, now);
  }
}
int climb(int x, int y, int nowcol) {
  while (x != y) {
    if (de[x] < de[y]) swap(x, y);
    if (col[fa[x]] == nowcol) return y;
    if (col[fa[x]]) return -1;
    col[fa[x]] = nowcol;
    x = fa[x];
  }
  return x;
}
void dsuDfs(int now, int last) {
  for (auto i : ng[now]) {
    if (i == last) continue;
    dsuDfs(i, now);
  }
  if (col[now]) {
    dp[now][1] = 1;
    for (auto i : ng[now]) {
      if (i == last) continue;
      dp[now][1] = dp[now][1] * ((dp[i][0] + dp[i][1]) % mod) % mod;
    }
  }
  if (col[now] == 0) {
    dp[now][0] = 1;
    for (auto i : ng[now]) {
      if (i == last) continue;
      dp[now][0] = dp[now][0] * ((dp[i][0] + dp[i][1]) % mod) % mod;
    }
    dp[now][1] = 0;
    have.clear();
    for (auto i : ng[now]) {
      if (i == last) continue;
      have.push_back(i);
    }
    pa[0] = 1, suf[(int)have.size() + 1] = 1;
    for (int i = 1; i <= (int)have.size(); i++) {
      pa[i] = pa[i - 1] * (dp[have[i - 1]][1] + dp[have[i - 1]][0]) % mod;
    }
    for (int i = (int)have.size(); i >= 1; i--) {
      suf[i] = suf[i + 1] * (dp[have[i - 1]][1] + dp[have[i - 1]][0]) % mod;
    }
    for (int i = 1; i <= (int)have.size(); i++) {
      long long Ad = pa[i - 1] * suf[i + 1] % mod * dp[have[i - 1]][1] % mod;
      dp[now][1] = (dp[now][1] + Ad) % mod;
    }
  }
}
int main() {
  int n, k;
  cin >> n >> k;
  for (int i = 1; i <= n; i++) scanf("%d", &col[i]), same[col[i]].push_back(i);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d%d", &u, &v);
    g[u].push_back(v), g[v].push_back(u);
  }
  de[1] = 0;
  fa[1] = -1;
  dfs(1, -1);
  for (int i = 1; i <= k; i++) {
    if ((int)same[i].size() == 0) continue;
    int h = same[i][0];
    for (int j = 1; j < (int)same[i].size(); j++) {
      h = climb(h, same[i][j], i);
      if (h == -1) {
        puts("0");
        return 0;
      }
    }
    tp[i] = h;
  }
  for (int i = 1; i <= n; i++) {
    if (!col[i])
      TP[i] = i;
    else
      TP[i] = tp[col[i]];
  }
  for (int i = 1; i <= n; i++) {
    for (auto j : g[i]) {
      if (TP[i] == TP[j]) continue;
      ng[TP[i]].push_back(TP[j]);
    }
  }
  dsuDfs(TP[1], -1);
  cout << dp[1][1] << '\n';
  return 0;
}
