#include <bits/stdc++.h>
using namespace std;
using ll = long long;
template <class T, class U>
ostream& operator<<(ostream& o, const pair<T, U>& p) {
  o << "(" << p.first << "," << p.second << ")";
  return o;
}
template <class T>
ostream& operator<<(ostream& o, const vector<T>& v) {
  o << "[";
  for (T t : v) {
    o << t << ",";
  }
  o << "]";
  return o;
}
const int MAX_V = 300003;
const int MAX_LOG_V = 19;
vector<int> G[MAX_V];
int root = 0;
int parent[MAX_LOG_V][MAX_V];
int depth[MAX_V];
void lca_dfs(int v, int p, int d) {
  parent[0][v] = p;
  depth[v] = d;
  for (int i = 0; i < (int)(G[v].size()); ++i) {
    if (G[v][i] != p) {
      lca_dfs(G[v][i], v, d + 1);
    }
  }
}
void lca_init(int V) {
  lca_dfs(root, -1, 0);
  for (int k = 0; k + 1 < MAX_LOG_V; ++k) {
    for (int v = 0; v < V; ++v) {
      if (parent[k][v] < 0)
        parent[k + 1][v] = -1;
      else
        parent[k + 1][v] = parent[k][parent[k][v]];
    }
  }
}
int lca(int u, int v) {
  if (depth[u] > depth[v]) swap(u, v);
  for (int k = 0; k < MAX_LOG_V; ++k) {
    if ((depth[v] - depth[u]) >> k & 1) v = parent[k][v];
  }
  if (u == v) return u;
  for (int k = MAX_LOG_V - 1; k >= 0; --k) {
    if (parent[k][u] != parent[k][v]) {
      u = parent[k][u];
      v = parent[k][v];
    }
  }
  return parent[0][u];
}
const int N = MAX_V;
const ll mod = 998244353;
int n, k;
int a[N];
int L[N];
bool invalid = false;
ll dp[N][2] = {};
int dfs(int x, int par) {
  map<int, int> m;
  for (int e : G[x])
    if (e != par) {
      int res = dfs(e, x);
      if (res != 0) m[e] = res;
    }
  if (m.empty()) {
    int idx = (a[x] != 0);
    dp[x][idx] = 1;
    if (L[a[x]] == x) return -1;
    return a[x];
  }
  unordered_set<int> s;
  vector<int> ch;
  for (const auto& p : m) {
    ch.push_back(p.first);
    if (p.second != -1) s.insert(p.second);
  }
  if (s.size() > 1) {
    invalid = true;
    return 0;
  }
  if (a[x] == 0 && s.size() == 0) {
    dp[x][0] = 1;
    for (int v : ch) (dp[x][0] *= (dp[v][0] + dp[v][1])) %= mod;
    int C = ch.size();
    vector<ll> pre(C), suf(C);
    for (int i = 0; i < (int)(C); ++i)
      pre[i] = suf[i] = (dp[ch[i]][0] + dp[ch[i]][1]) % mod;
    for (int i = 0; i < (int)(C - 1); ++i) (pre[i + 1] *= pre[i]) %= mod;
    for (int i = C - 1; i > 0; --i) (suf[i - 1] *= suf[i]) %= mod;
    dp[x][1] = 0;
    for (int i = 0; i < (int)(C); ++i) {
      ll mul = 1;
      if (i != 0) (mul *= pre[i - 1]) %= mod;
      if (i != C - 1) (mul *= suf[i + 1]) %= mod;
      ll add = (mul * dp[ch[i]][1]) % mod;
      (dp[x][1] += add) %= mod;
    }
    return -1;
  } else {
    int col = a[x];
    if (!s.empty()) {
      assert(s.size() == 1);
      int h = *s.begin();
      if (col != 0 && col != h) {
        invalid = true;
        return 0;
      }
      col = h;
    }
    dp[x][1] = 1;
    for (int v : ch) {
      if (m[v] == -1)
        (dp[x][1] *= (dp[v][0] + dp[v][1])) %= mod;
      else
        (dp[x][1] *= dp[v][1]) %= mod;
    }
    return (x == L[col]) ? -1 : col;
  }
}
int main() {
  scanf(" %d %d", &n, &k);
  vector<vector<int>> col(k + 1);
  for (int i = 0; i < (int)(n); ++i) {
    scanf(" %d", &a[i]);
    col[a[i]].push_back(i);
  }
  for (int i = 0; i < (int)(n - 1); ++i) {
    int u, v;
    scanf(" %d %d", &u, &v);
    --u;
    --v;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  lca_init(n);
  for (int i = 1; i <= k; ++i) {
    int x = col[i][0];
    for (int j : col[i]) x = lca(x, j);
    if (a[x] != 0 && a[x] != i) invalid = true;
    a[x] = i;
    L[i] = x;
  }
  dfs(0, -1);
  printf("%lld\n", invalid ? 0 : dp[0][1]);
  return 0;
}
