#include <bits/stdc++.h>
using namespace std;
int n, k, a[300005], p[300005], l[300005], r[300005], deep[300005], vis[300005],
    fa[300005], flag;
long long dp[300005][3];
vector<int> G[300005], V[300005], E[300005];
void dfs(int x, int f) {
  for (auto u : G[x])
    if (u != f) {
      fa[u] = x;
      deep[u] = deep[x] + 1;
      dfs(u, x);
    }
}
void work(int x, int f) {
  dp[x][0] = 1;
  for (auto u : E[x])
    if (u != f) {
      work(u, x);
      dp[x][1] = (dp[x][1] * (dp[u][1] + dp[u][0]) % 998244353 +
                  dp[x][0] * dp[u][1] % 998244353) %
                 998244353;
      dp[x][0] = (dp[x][0] * (dp[u][1] + dp[u][0]) % 998244353) % 998244353;
    }
  if (a[x]) dp[x][1] = dp[x][0], dp[x][0] = 0;
}
int main() {
  scanf("%d%d", &n, &k);
  for (int i = 1; i <= n; i++) {
    scanf("%d", &a[i]);
    V[a[i]].push_back(i);
  }
  for (int i = 1; i < n; i++) {
    scanf("%d%d", &l[i], &r[i]);
    G[l[i]].push_back(r[i]);
    G[r[i]].push_back(l[i]);
  }
  dfs(1, -1);
  flag = 0;
  for (int i = 1; i <= k; i++) {
    int now = V[i][0];
    vis[now] = 1;
    for (int j = 1; j < V[i].size(); j++)
      if (!vis[V[i][j]]) {
        int nx = V[i][j];
        int Flag = 0;
        vis[nx] = 1;
        if (deep[now] > deep[nx]) swap(nx, now);
        while (deep[nx] != deep[now]) {
          nx = fa[nx];
          if (a[nx] != a[now] && a[nx] != 0) {
            flag = 1;
            break;
          }
          if (vis[nx]) {
            Flag = 1;
            break;
          }
          a[nx] = a[now];
          vis[nx] = 1;
        }
        if (flag) break;
        if (Flag) continue;
        while (nx != now) {
          now = fa[now];
          nx = fa[nx];
          vis[now] = vis[nx] = 1;
          if (a[now] != 0 && a[now] != i) {
            flag = 1;
            break;
          }
          if (a[nx] != 0 && a[nx] != i) {
            flag = 1;
            break;
          }
          a[now] = a[nx] = i;
        }
        if (flag) break;
      }
    if (flag) break;
  }
  if (flag) return printf("0\n"), 0;
  for (int i = 1; i <= n; i++)
    if (!p[a[i]]) p[a[i]] = i;
  for (int i = 1; i < n; i++) {
    int u = l[i], v = r[i];
    if (a[u]) u = p[a[u]];
    if (a[v]) v = p[a[v]];
    if (u != v) E[v].push_back(u), E[u].push_back(v);
  }
  work(1, -1);
  printf("%lld\n", dp[1][1]);
}
