#include <bits/stdc++.h>
using namespace std;
int n, K, a[600005], cnt, head[600005], tot, pos[600005], far[600005];
int deep[600005], size[600005], fa[600005], son[600005], top[600005];
int num, p[600005 << 1], topp, s[600005];
bool flag;
vector<int> col[600005];
struct edge {
  int u, v;
} E[600005];
struct node {
  int to, next;
} e[600005 << 1];
void insert(int u, int v) {
  e[++cnt].to = v;
  e[cnt].next = head[u];
  head[u] = cnt;
  e[++cnt].to = u;
  e[cnt].next = head[v];
  head[v] = cnt;
}
void dfs1(int x) {
  size[x] = 1;
  son[x] = 0;
  for (int i = head[x]; i; i = e[i].next)
    if (e[i].to ^ fa[x]) {
      deep[e[i].to] = deep[x] + 1;
      fa[e[i].to] = x;
      dfs1(e[i].to);
      size[x] += size[e[i].to];
      if (size[son[x]] < size[e[i].to]) son[x] = e[i].to;
    }
}
void dfs2(int x, int chain) {
  pos[x] = ++tot;
  top[x] = chain;
  if (!son[x]) {
    far[x] = tot;
    return;
  }
  dfs2(son[x], chain);
  for (int i = head[x]; i; i = e[i].next)
    if (e[i].to ^ fa[x] && e[i].to ^ son[x]) dfs2(e[i].to, e[i].to);
  far[x] = tot;
}
int LCA(int x, int y) {
  while (top[x] ^ top[y]) {
    if (deep[top[x]] < deep[top[y]]) swap(x, y);
    x = fa[top[x]];
  }
  return deep[x] < deep[y] ? x : y;
}
class Vitual_T {
 public:
  int cnt, head[600005], F[600005][2];
  struct node {
    int to, next, v;
  } e[600005 << 1];
  void insert(int u, int v, int w) {
    e[++cnt].to = v;
    e[cnt].next = head[u];
    e[cnt].v = w;
    head[u] = cnt;
  }
  int add(int x, int y) {
    x += y;
    if (x >= 998244353) x -= 998244353;
    return x;
  }
  int Pow(int x, int p) {
    int t = 1;
    for (; p; p >>= 1, x = (long long)x * x % 998244353)
      if (p & 1) t = (long long)t * x % 998244353;
    return t;
  }
  int inv(int x) { return Pow(x, 998244353 - 2); }
  void dfs(int x) {
    int mul = 1;
    for (int i = head[x]; i; i = e[i].next) {
      dfs(e[i].to);
      mul = (long long)mul *
            add(F[e[i].to][0], (long long)F[e[i].to][1] * e[i].v % 998244353) %
            998244353;
    }
    if (x <= n) {
      F[x][0] = mul;
      for (int i = head[x]; i; i = e[i].next) {
        int y = e[i].to;
        F[x][1] =
            add(F[x][1],
                (long long)F[y][1] * mul % 998244353 *
                    inv(add(F[y][0], (long long)F[y][1] * e[i].v % 998244353)) %
                    998244353);
      }
    } else {
      F[x][1] = mul;
    }
  }
  void solve() {
    dfs(n + 1);
    printf("%d\n", F[n + 1][1]);
  }
} T;
bool cmp(int x, int y) { return pos[x] < pos[y]; }
bool con(int x, int y) { return pos[x] <= pos[y] && far[x] >= pos[y]; }
void color(int x, int y, int c) {
  while (y != fa[x]) {
    if (a[y]) {
      if (a[y] != c) {
        flag = 1;
        return;
      }
    } else
      a[y] = c;
    y = fa[y];
  }
}
void build() {
  sort(p + 1, p + 1 + num, cmp);
  int pre = num;
  for (int i = 1; i < pre; i++) p[++num] = LCA(p[i], p[i + 1]);
  sort(p + 1, p + 1 + num, cmp);
}
void work(int x) {
  num = 0;
  for (int i = 0; i < col[x].size(); i++) p[++num] = col[x][i];
  build();
  topp = 0;
  s[++topp] = p[1];
  for (int i = 2; i <= num; i++) {
    if (p[i] == s[topp]) continue;
    while (topp > 0 && !con(s[topp], p[i])) topp--;
    color(s[topp], p[i], x);
    if (flag) return;
    s[++topp] = p[i];
  }
}
void rebuild() {
  num = 0;
  for (int i = 1; i <= K; i++) p[++num] = n + i;
  build();
  topp = 0;
  s[++topp] = p[1];
  for (int i = 2; i <= num; i++) {
    if (p[i] == s[topp]) continue;
    while (topp > 0 && !con(s[topp], p[i])) topp--;
    T.insert(s[topp], p[i], deep[p[i]] - deep[s[topp]]);
    s[++topp] = p[i];
  }
  T.solve();
}
int main() {
  scanf("%d%d", &n, &K);
  for (int i = 1; i <= n; i++) {
    scanf("%d", &a[i]);
    if (a[i] > 0) col[a[i]].push_back(i);
  }
  for (int i = 1; i < n; i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    insert(x, y);
    E[i].u = x;
    E[i].v = y;
  }
  dfs1(1);
  dfs2(1, 1);
  for (int i = 1; i <= K; i++) {
    work(i);
    if (flag) {
      puts("0");
      return 0;
    }
  }
  cnt = 0;
  for (int i = 1; i <= n; i++) head[i] = 0;
  for (int i = 1; i < n; i++) {
    int x, y;
    if (a[E[i].u])
      x = n + a[E[i].u];
    else
      x = E[i].u;
    if (a[E[i].v])
      y = n + a[E[i].v];
    else
      y = E[i].v;
    if (x != y) insert(x, y);
  }
  tot = 0;
  dfs1(n + 1);
  dfs2(n + 1, n + 1);
  rebuild();
  return 0;
}
