#include <bits/stdc++.h>
using namespace std;
char getc() {
  char c = getchar();
  while ((c < 'A' || c > 'Z') && (c < 'a' || c > 'z') && (c < '0' || c > '9'))
    c = getchar();
  return c;
}
int gcd(int n, int m) { return m == 0 ? n : gcd(m, n % m); }
int read() {
  int x = 0, f = 1;
  char c = getchar();
  while (c < '0' || c > '9') {
    if (c == '-') f = -1;
    c = getchar();
  }
  while (c >= '0' && c <= '9')
    x = (x << 1) + (x << 3) + (c ^ 48), c = getchar();
  return x * f;
}
int n, m, a[300010], p[300010 << 1], dfn[300010], fa[300010][20], deep[300010],
    f[300010 << 1], F[300010 << 1][2], pre[300010 << 1], suf[300010 << 1], t, T,
    cnt, stk[300010], top;
vector<int> pos[300010];
struct data {
  int to, nxt;
} edge[300010 << 1], tree[300010 << 2];
void addedge(int x, int y) {
  t++;
  edge[t].to = y, edge[t].nxt = p[x], p[x] = t;
}
void new_addedge(int x, int y) {
  T++;
  tree[T].to = y, tree[T].nxt = p[x], p[x] = T;
}
void dfs(int k) {
  dfn[k] = ++cnt;
  for (int i = p[k]; i; i = edge[i].nxt)
    if (edge[i].to != fa[k][0]) {
      fa[edge[i].to][0] = k;
      deep[edge[i].to] = deep[k] + 1;
      dfs(edge[i].to);
    }
}
bool cmp(const int& a, const int& b) { return dfn[a] < dfn[b]; }
int lca(int x, int y) {
  if (deep[x] < deep[y]) swap(x, y);
  for (int j = 19; ~j; j--)
    if (deep[fa[x][j]] >= deep[y]) x = fa[x][j];
  if (x == y) return x;
  for (int j = 19; ~j; j--)
    if (fa[x][j] != fa[y][j]) x = fa[x][j], y = fa[y][j];
  return fa[x][0];
}
bool paint(int x, int y, int color) {
  while (x != y) {
    if (f[x] > n && f[x] != color) return 1;
    f[x] = color;
    x = fa[x][0];
  }
  return 0;
}
void dp(int k, int from) {
  for (int i = p[k]; i; i = tree[i].nxt)
    if (tree[i].to != from) dp(tree[i].to, k);
  if (k > n) {
    F[k][0] = 0;
    F[k][1] = 1;
    for (int i = p[k]; i; i = tree[i].nxt)
      if (tree[i].to != from)
        F[k][1] =
            1ll * F[k][1] * (F[tree[i].to][0] + F[tree[i].to][1]) % 998244353;
  } else {
    F[k][0] = 1;
    int cnt = 0;
    for (int i = p[k]; i; i = tree[i].nxt)
      if (tree[i].to != from) {
        F[k][0] =
            1ll * F[k][0] * (F[tree[i].to][0] + F[tree[i].to][1]) % 998244353;
        pre[++cnt] = F[tree[i].to][0] + F[tree[i].to][1];
      }
    for (int i = 1; i <= cnt; i++) suf[i] = pre[i];
    pre[0] = 1;
    for (int i = 1; i <= cnt; i++)
      pre[i] = 1ll * pre[i - 1] * pre[i] % 998244353;
    suf[cnt + 1] = 1;
    for (int i = cnt; i >= 1; i--)
      suf[i] = 1ll * suf[i] * suf[i + 1] % 998244353;
    int t = 0;
    for (int i = p[k]; i; i = tree[i].nxt)
      if (tree[i].to != from) {
        t++;
        F[k][1] = (F[k][1] + 1ll * pre[t - 1] * suf[t + 1] % 998244353 *
                                 F[tree[i].to][1]) %
                  998244353;
      }
  }
}
signed main() {
  n = read(), m = read();
  for (int i = 1; i <= n; i++) a[i] = read();
  for (int i = 1; i < n; i++) {
    int x = read(), y = read();
    addedge(x, y), addedge(y, x);
  }
  fa[1][0] = 1;
  dfs(1);
  for (int j = 1; j < 20; j++)
    for (int i = 1; i <= n; i++) fa[i][j] = fa[fa[i][j - 1]][j - 1];
  for (int i = 1; i <= n; i++)
    if (a[i])
      f[i] = n + a[i], pos[a[i]].push_back(i);
    else
      f[i] = i;
  for (int i = 1; i <= m; i++)
    if (pos[i].size() > 1) {
      sort(pos[i].begin(), pos[i].end(), cmp);
      int root = (*pos[i].begin());
      for (int j = 0; j < pos[i].size() - 1; j++)
        if (deep[lca(pos[i][j], pos[i][j + 1])] < deep[root])
          root = lca(pos[i][j], pos[i][j + 1]);
      if (f[root] > n && f[root] != n + i) {
        cout << 0;
        return 0;
      }
      top = 0;
      stk[++top] = root;
      f[root] = n + i;
      for (int j = ((*pos[i].begin()) == root); j < pos[i].size(); j++) {
        int l = lca(stk[top], pos[i][j]);
        if (l != stk[top]) {
          while (top > 1 && deep[l] <= deep[stk[top - 1]]) {
            if (paint(stk[top], stk[top - 1], n + i)) {
              cout << 0;
              return 0;
            }
            top--;
          }
          if (paint(stk[top], l, n + i)) {
            cout << 0;
            return 0;
          }
          stk[top] = l;
        }
        stk[++top] = pos[i][j];
      }
      while (top > 1) {
        if (paint(stk[top], stk[top - 1], n + i)) {
          cout << 0;
          return 0;
        }
        top--;
      }
    }
  memset(p, 0, sizeof(p));
  for (int i = 1; i <= t; i += 2)
    if (f[edge[i].to] != f[edge[i + 1].to])
      new_addedge(f[edge[i].to], f[edge[i + 1].to]),
          new_addedge(f[edge[i + 1].to], f[edge[i].to]);
  dp(n + 1, n + 1);
  cout << F[n + 1][1];
  return 0;
}
