#include <bits/stdc++.h>
using namespace std;
const long long MOD = 998244353LL;
const long double PI = 3.141592653589793238462643383279502884197;
priority_queue<int, vector<int>, greater<int> > pq;
vector<int> col[300001], v[300000], vv[300000];
queue<int> q, use;
int n, k;
int vec[300000];
int p[300000];
int cc[300000];
int dist[300000];
int spa[300000][19];
long long dp[300000][2];
char visit[300000];
long long mp(long long a, long long b) {
  long long ret = 1;
  while (b) {
    if (b % 2) ret = ret * a % MOD;
    a = a * a % MOD;
    b /= 2;
  }
  return ret;
}
void go1(int o, int pa, int di) {
  dist[o] = di;
  if (o) {
    spa[o][0] = pa;
    for (int i = 1; i < 19; i++) spa[o][i] = spa[spa[o][i - 1]][i - 1];
  }
  for (int i = 0; i < (int)(v[o]).size(); i++) {
    if (v[o][i] == pa) continue;
    go1(v[o][i], o, di + 1);
  }
}
int lca(int a, int b) {
  if (dist[a] > dist[b]) swap(a, b);
  for (int i = 18; i >= 0; i--)
    if (dist[spa[b][i]] >= dist[a]) b = spa[b][i];
  if (a == b) return a;
  for (int i = 18; i >= 0; i--)
    if (spa[a][i] != spa[b][i]) {
      a = spa[a][i];
      b = spa[b][i];
    }
  return spa[a][0];
}
int find(int a) {
  if (p[a] == -1) return a;
  return p[a] = find(p[a]);
}
void merge(int a, int b) {
  a = find(a);
  b = find(b);
  if (a == b) return;
  if (dist[a] > dist[b]) swap(a, b);
  if (!cc[a]) cc[a] = cc[b];
  if (!cc[b]) cc[b] = cc[a];
  p[b] = a;
}
void con(int o, int rp) {
  visit[o] = 1;
  for (int i = 0; i < (int)(v[o]).size(); i++) {
    if (find(v[o][i]) != rp)
      vv[rp].push_back(find(v[o][i]));
    else if (dist[o] < dist[v[o][i]])
      con(v[o][i], rp);
  }
}
long long go(int o, int pa, int is) {
  if (cc[o] && is) return dp[o][is] = go(o, pa, 0);
  if (dp[o][is] != -1) return dp[o][is];
  long long ret = 1, tmp = 1;
  if (is == 0) {
    for (int i = 0; i < (int)(vv[o]).size(); i++) {
      if (vv[o][i] == pa) continue;
      if (cc[vv[o][i]])
        tmp = go(vv[o][i], o, 0);
      else
        tmp = go(vv[o][i], o, 0) + go(vv[o][i], o, 1);
      if (tmp) ret = ret * tmp % MOD;
    }
  } else {
    ret = 0;
    tmp = go(o, pa, 0);
    for (int i = 0; i < (int)(vv[o]).size(); i++) {
      if (vv[o][i] == pa) continue;
      long long tt = go(vv[o][i], o, 0), ta = go(vv[o][i], o, 1);
      if (!tt || !ta) continue;
      if (cc[vv[o][i]])
        ret = (ret + tmp) % MOD;
      else
        ret = (ret + tmp * mp(tt + ta, MOD - 2) % MOD * ta) % MOD;
    }
  }
  if ((int)(vv[o]).size() == 1 && cc[o] == 0) return dp[o][is] = 0;
  return dp[o][is] = ret;
}
int main() {
  memset(p, -1, sizeof(p));
  memset(dp, -1, sizeof(dp));
  scanf("%d %d", &n, &k);
  for (int i = 0; i < n; i++) {
    scanf("%d", &vec[i]);
    cc[i] = vec[i];
    if (vec[i]) col[vec[i]].push_back(i);
  }
  int a, b;
  for (int i = 0; i < n - 1; i++) {
    scanf("%d %d", &a, &b);
    a--;
    b--;
    v[a].push_back(b);
    v[b].push_back(a);
  }
  go1(0, 0, 0);
  for (int i = 1; i <= k; i++) {
    int lc = col[i][0];
    q.push(col[i][0]);
    visit[col[i][0]] = 1;
    for (int j = 1; j < (int)(col[i]).size(); j++) {
      lc = lca(lc, col[i][j]);
      q.push(col[i][j]);
      visit[col[i][j]] = 1;
    }
    int tv;
    while (!q.empty()) {
      tv = q.front();
      q.pop();
      use.push(tv);
      if (tv == lc) continue;
      int ob = spa[tv][0];
      if (vec[ob] != 0 && vec[tv] != 0 && vec[ob] != vec[tv])
        return !printf("0");
      merge(ob, tv);
      if (!visit[ob]) {
        q.push(ob);
        visit[ob] = 1;
      }
    }
    while (!use.empty()) {
      visit[use.front()] = 0;
      use.pop();
    }
  }
  for (int i = 0; i < n; i++)
    if (!visit[find(i)]) con(find(i), find(i));
  printf("%lld", go(find(col[1][0]), find(col[1][0]), 0));
}
