#include <bits/stdc++.h>
using namespace std;
int pct(int x) { return __builtin_popcount(x); }
int pct(long long x) { return __builtin_popcountll(x); }
int bit(int x) { return 31 - __builtin_clz(x); }
int bit(long long x) { return 63 - __builtin_clzll(x); }
int cdiv(int a, int b) { return a / b + !(a < 0 || a % b == 0); }
long long cdiv(long long a, long long b) {
  return a / b + !(a < 0 || a % b == 0);
}
template <typename T>
void leftShift(vector<T> &v, long long k) {
  k %= (int)v.size();
  if (k < 0) k += (int)v.size();
  rotate(v.begin(), v.begin() + k, v.end());
}
template <typename T>
void rightShift(vector<T> &v, long long k) {
  leftShift(v, (int)v.size() - k);
}
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
long long rand(long long l, long long r) {
  uniform_int_distribution<long long> uid(l, r);
  return uid(rng);
}
void pr() {}
void sc() {}
template <typename Head, typename... Tail>
void pr(Head H, Tail... T) {
  cout << H << " ";
  pr(T...);
}
template <typename Head, typename... Tail>
void sc(Head &H, Tail &...T) {
  cin >> H;
  sc(T...);
}
string to_string(__int128 x) {
  string s = "";
  bool neg = 0;
  if (x < 0) {
    s += "-";
    neg = 1;
    x = -x;
  }
  if (!x) s += '0';
  while (x) {
    int rem = x % 10;
    s += to_string(rem);
    x /= 10;
  }
  reverse(s.begin() + neg, s.end());
  return s;
}
const int mod = 998244353;
int pwr(int a, int b) {
  int ans = 1;
  while (b) {
    if (b & 1) ans = (ans * 1LL * a) % mod;
    a = (a * 1LL * a) % mod;
    b >>= 1;
  }
  return ans;
}
const int N = 3e5 + 5;
int a[N];
vector<int> g[N];
int depth[N];
int dsu[N];
int par[N];
int T[N], tt;
int root(int u) { return ((dsu[u] == u) ? u : dsu[u] = root(dsu[u])); }
void uni(int u, int p) {
  u = root(u), p = root(p);
  dsu[u] = p;
}
vector<int> cnodes[N];
void pre_dfs(int u, int p) {
  T[u] = ++tt;
  par[u] = p;
  depth[u] = depth[p] + 1;
  for (int v : g[u]) {
    if (v != p) {
      if (a[v] && a[v] == a[u]) {
        uni(v, u);
      }
      pre_dfs(v, u);
    }
  }
}
void color(int u, int v, int c) {
  while (u != v) {
    if (depth[u] < depth[v]) swap(u, v);
    int p = root(par[u]);
    if (!a[p])
      a[p] = a[u];
    else if (a[p] != a[u]) {
      cout << 0;
      exit(0);
    }
    uni(u, p);
    u = p;
  }
}
int ans = 1;
array<int, 2> dp[N];
array<int, 2> dp1[2];
array<int, 2> dp2[N][2];
void go(int u, int p) {
  if (a[u]) {
    dp[u][0] = 0;
    dp[u][1] = 1;
    return;
  }
  dp2[u][0][0] = 1;
  dp2[u][0][1] = 1;
  for (int v : g[u]) {
    if (v != p) {
      go(v, u);
      memset(dp1, 0, sizeof dp1);
      for (int i = 0; i <= 1; i++) {
        for (int j = 0; j <= 1; j++) {
          int val = dp2[u][i][j];
          for (int k = 0; k <= 1; k++) {
            int val1 = dp[v][k];
            if ((i + k) < 2) {
              dp1[i | k][j] = (dp1[i | k][j] + val * 1LL * val1) % mod;
            }
            if (k) {
              dp1[i][j] = (dp1[i][j] + val * 1LL * val1) % mod;
            }
          }
        }
      }
      for (int i = 0; i <= 1; i++) {
        for (int j = 0; j <= 1; j++) {
          dp2[u][i][j] = dp1[i][j];
        }
      }
    }
  }
  for (int i = 0; i <= 1; i++) {
    dp[u][i] = dp2[u][i][i];
  }
}
void dfs(int u, int p) {
  if (!a[u] && a[p]) {
    go(u, p);
    ans = (ans * 1LL * (dp[u][0] + dp[u][1])) % mod;
  }
  for (int v : g[u]) {
    if (v != p) {
      dfs(v, u);
    }
  }
}
void solve() {
  int n, k;
  sc(n, k);
  for (int i = 1; i <= n; i++) {
    dsu[i] = i;
    sc(a[i]);
    if (a[i]) cnodes[a[i]].push_back(i);
  }
  for (int i = 1; i <= n - 1; i++) {
    int u, v;
    sc(u, v);
    g[u].push_back(v);
    g[v].push_back(u);
  }
  int root = 1;
  for (int i = 1; i <= n; i++) {
    if (a[i]) {
      root = i;
      break;
    }
  }
  pre_dfs(root, 0);
  auto cmp = [&](const int &a, const int &b) { return T[a] < T[b]; };
  for (int i = 1; i <= k; i++) {
    sort(cnodes[i].begin(), cnodes[i].end());
    for (int j = 0; j <= (int)cnodes[i].size() - 2; j++) {
      color(cnodes[i][j], cnodes[i][j + 1], i);
    }
  }
  dfs(root, 0);
  cout << ans;
}
signed main() {
  ios ::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  int t;
  t = 1;
  for (int tt = 1; tt <= t; tt++) {
    solve();
  }
  return 0;
}
