#include <bits/stdc++.h>
using namespace std;
const int N = 3e5 + 1, logN = 20 + 1, mod = 998244353;
int n, k;
vector<int> g[N];
int c[N];
vector<int> rc[N];
int clcm[N];
int t;
int h[N], p[N];
int tin[N], tout[N];
int order[2 * N];
void euler_tour(int v = 0, int par = -1, int hgt = 0) {
  order[t] = v;
  tin[v] = t++;
  h[v] = hgt;
  p[v] = par;
  for (int u : g[v])
    if (u != par) {
      euler_tour(u, v, hgt + 1);
      order[t++] = v;
    }
  tout[v] = t;
}
int fl(int a) { return 31 - __builtin_clz(a); }
pair<int, int> mintable[logN][2 * N];
void init_table() {
  for (int i = 0; i < 2 * n - 1; ++i) mintable[0][i] = {h[order[i]], order[i]};
  for (int level = 0; level < logN - 1; ++level)
    for (int i = 0; i <= 2 * n - 1 - (1 << level + 1); ++i)
      mintable[level + 1][i] =
          min(mintable[level][i], mintable[level][i + (1 << level)]);
}
int get(int l, int r) {
  int level = fl(r - l);
  return min(mintable[level][l], mintable[level][r - (1 << level)]).second;
}
int lcm(int v, int u) {
  return get(min(tin[v], tin[u]), max(tout[v], tout[u]));
}
void paint_dfs(int v, int top, int clr) {
  if (c[v] == clr) return;
  if (c[v] != 0) {
    cout << "0\n";
    exit(0);
  }
  c[v] = clr;
  if (v == top) return;
  paint_dfs(p[v], top, clr);
}
int new_v[N];
int sz;
vector<int> new_g[N];
void new_dfs(int v, int p = -1) {
  for (int u : g[v]) {
    if (u == p) continue;
    if (new_v[v] != new_v[u]) {
      new_g[new_v[v]].push_back(new_v[u]);
      new_g[new_v[u]].push_back(new_v[v]);
    }
    new_dfs(u, v);
  }
}
long long f[2][N];
void dfs(int v, int p = -1) {
  f[0][v] = 0;
  f[1][v] = 1;
  for (int u : new_g[v]) {
    if (u == p) continue;
    dfs(u, v);
    f[0][v] = (f[0][v] * f[1][u] + f[1][v] * f[0][u]) % mod;
    f[1][v] = (f[1][v] * f[1][u]) % mod;
  }
  if (v < k)
    f[0][v] = f[1][v];
  else
    f[1][v] = (f[1][v] + f[0][v]) % mod;
}
int main() {
  ios_base::sync_with_stdio(false);
  cout.tie(0);
  cin.tie(0);
  cin >> n >> k;
  for (int i = 0; i < n; ++i) {
    cin >> c[i];
    rc[c[i]].push_back(i);
  }
  for (int i = 0; i < n - 1; ++i) {
    int u, v;
    cin >> u >> v;
    --u;
    --v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  euler_tour();
  init_table();
  for (int i = 1; i <= k; ++i) {
    clcm[i] = rc[i][0];
    for (int j = 1; j < rc[i].size(); ++j) clcm[i] = lcm(clcm[i], rc[i][j]);
    for (int v : rc[i])
      if (v != clcm[i]) paint_dfs(p[v], clcm[i], i);
  }
  sz = k;
  for (int i = 0; i < n; ++i) new_v[i] = (c[i] ? c[i] - 1 : sz++);
  new_dfs(0);
  n = sz;
  dfs(0);
  cout << f[0][0] << '\n';
}
