#include <bits/stdc++.h>
using namespace std;
const long long MXN = 3e5 + 10;
const long long Mod = 998244353;
const long long LOG = 20;
long long n, k;
long long A[MXN], has[MXN], dp[2][MXN];
long long Par[MXN], Sz[MXN], R[2][MXN];
long long Jad[LOG][MXN], dis[MXN];
vector<long long> adj[MXN], G[MXN], Col[MXN];
vector<pair<long long, long long>> E;
inline void mkay(long long& x) {
  if (x >= Mod) x -= Mod;
}
long long power(long long a, long long b) {
  return (!b ? 1 : power(a * a % Mod, b / 2) * (b & 1LL ? a : 1) % Mod);
}
long long Find(long long x) {
  return (x == Par[x] ? x : Par[x] = Find(Par[x]));
}
void Union(long long x, long long y) {
  x = Find(x), y = Find(y);
  if (x == y) return;
  if (Sz[x] < Sz[y]) swap(x, y);
  Sz[x] += Sz[y], Par[y] = x;
  if (R[0][x] == 0)
    R[0][x] = R[0][y], R[1][x] = R[1][y];
  else if (R[0][y] > 0) {
    R[0][x] = min(R[0][x], R[0][y]);
    R[1][x] = max(R[1][x], R[1][y]);
  }
}
void prep(long long u, long long par) {
  Jad[0][u] = par;
  if (A[u]) Col[A[u]].push_back(u);
  for (int i = 1; i < LOG; i++) {
    Jad[i][u] = Jad[i - 1][Jad[i - 1][u]];
  }
  for (auto v : adj[u]) {
    if (v != par) dis[v] = dis[u] + 1, prep(v, u);
  }
}
long long K_Jad(long long u, long long k) {
  for (int i = 0; i < LOG; i++) {
    if ((k >> i) & 1LL) u = Jad[i][u];
  }
  return u;
}
long long LCA(long long u, long long v) {
  if (dis[v] < dis[u]) swap(u, v);
  v = K_Jad(v, dis[v] - dis[u]);
  if (u == v) return u;
  for (int i = LOG - 1; ~i; i--) {
    if (Jad[i][u] != Jad[i][v]) u = Jad[i][u], v = Jad[i][v];
  }
  return Jad[0][u];
}
void flood(long long u, long long par) {
  for (auto v : adj[u]) {
    if (v == par) continue;
    flood(v, u), has[u] += has[v];
    if (has[v]) Union(u, v);
  }
}
void DFS(long long u, long long par) {
  dp[0][u] = 1;
  for (auto v : G[u]) {
    if (v == par) continue;
    DFS(v, u);
    dp[0][u] = dp[0][u] * ((dp[0][v] + dp[1][v]) % Mod) % Mod;
  }
  for (auto v : G[u]) {
    if (v == par) continue;
    long long other =
        dp[0][u] * power((dp[0][v] + dp[1][v]) % Mod, Mod - 2) % Mod;
    dp[1][u] += other * dp[1][v] % Mod;
    mkay(dp[1][u]);
  }
  if (R[0][u]) dp[1][u] = dp[0][u], dp[0][u] = 0;
}
int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cin >> n >> k;
  iota(Par + 1, Par + n + 1, 1), fill(Sz + 1, Sz + n + 1, 1);
  for (int i = 1; i <= n; i++) cin >> A[i], R[0][i] = R[1][i] = A[i];
  for (int i = 1; i < n; i++) {
    long long u, v;
    cin >> u >> v;
    adj[u].push_back(v), adj[v].push_back(u);
    E.push_back({u, v});
  }
  prep(1, 0);
  for (int i = 1; i <= k; i++) {
    for (int j = 1; j < Col[i].size(); j++) {
      long long u = Col[i][j - 1], v = Col[i][j];
      has[u]++, has[v]++, has[LCA(u, v)] -= 2;
    }
  }
  flood(1, 0);
  for (int i = 1; i <= n; i++) {
    if (i == Par[i]) {
      if (R[0][i] != R[1][i]) return cout << 0, 0;
    }
  }
  for (auto e : E) {
    long long u, v;
    tie(u, v) = e;
    u = Find(u), v = Find(v);
    if (u == v) continue;
    G[u].push_back(v), G[v].push_back(u);
  }
  DFS(Find(1), 0);
  cout << dp[1][Find(1)] << '\n';
  return 0;
}
