#include <bits/stdc++.h>
using namespace std;
const int64_t MOD = 998244353;
void add(int64_t& a, int64_t b) { a = (a + b) % MOD; }
void mul(int64_t& a, int64_t b) { a = a * b % MOD; }
struct LCA {
  static const int BITLEN_MAX = 30;
  vector<int> parent[BITLEN_MAX];
  vector<int> depth;
  int bitlen;
  void initialize(int N, const vector<int> edges[]) {
    int root = 0;
    bitlen = 1;
    while ((1 << bitlen) < N) bitlen += 1;
    for (int i = 0; i < bitlen; i++) parent[i].resize(N);
    depth.resize(N, -1);
    dfs(root, -1, 0, edges);
    for (int k = 0; k < bitlen - 1; k++) {
      for (int v = 0; v < N; v++) {
        if (depth[v] == -1) continue;
        if (parent[k][v] < 0) {
          parent[k + 1][v] = -1;
        } else {
          parent[k + 1][v] = parent[k][parent[k][v]];
        }
      }
    }
  }
  void dfs(int v, int p, int d, const vector<int> edges[]) {
    parent[0][v] = p;
    depth[v] = d;
    for (auto u : edges[v]) {
      if (u != p) dfs(u, v, d + 1, edges);
    }
  }
  int calc_lca(int u, int v) {
    if (depth[u] > depth[v]) swap(u, v);
    for (int k = 0; k < bitlen; k++) {
      if (((depth[v] - depth[u]) >> k) & 1) v = parent[k][v];
    }
    if (u == v) return u;
    for (int k = bitlen - 1; k >= 0; k--) {
      if (parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }
    return parent[0][u];
  }
  int calc_dist(int u, int v) {
    int l = calc_lca(u, v);
    return depth[u] + depth[v] - depth[l] * 2;
  }
};
int N, K, A[300000];
vector<int> edges[300000];
int64_t dp[300000][2];
void dfs2(int i, int p) {
  if (A[i]) {
    dp[i][1] = 1;
    for (int j : edges[i])
      if (j != p) {
        dfs2(j, i);
        if (A[j] == A[i]) {
          mul(dp[i][1], dp[j][1]);
        } else {
          mul(dp[i][1], dp[j][0] + dp[j][1]);
        }
      }
  } else {
    dp[i][0] = 1;
    vector<int64_t> dp2 = {1};
    vector<int> vs = {1};
    for (int j : edges[i])
      if (j != p) {
        dfs2(j, i);
        mul(dp[i][0], dp[j][0] + dp[j][1]);
        vs.push_back(j);
        dp2.push_back((dp[j][0] + dp[j][1]) % MOD);
      }
    dp2.push_back(1);
    vs.push_back(1);
    int sz = vs.size();
    vector<int64_t> ml(sz, 1), mr(sz, 1);
    for (int k = 1; k < sz - 1; k++) ml[k] = ml[k - 1] * dp2[k] % MOD;
    for (int k = sz - 2; k > 0; k--) mr[k] = mr[k + 1] * dp2[k] % MOD;
    for (int k = 1; k < sz - 1; k++) {
      int j = vs[k];
      int64_t result = dp[j][1] * ml[k - 1] % MOD * mr[k + 1] % MOD;
      add(dp[i][1], result);
    }
  }
}
vector<int> CV[300001];
bitset<300000> done;
int main() {
  cin >> N >> K;
  for (int i = 0; i < N; i++) {
    scanf("%d", &A[i]);
    CV[A[i]].push_back(i);
  }
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    edges[a - 1].push_back(b - 1);
    edges[b - 1].push_back(a - 1);
  }
  LCA lca;
  lca.initialize(N, edges);
  for (int i = 1; i <= K; i++) {
    if (CV[i].size() == 0) continue;
    int v = CV[i][0];
    for (int j = 1; j < CV[i].size(); j++) v = lca.calc_lca(v, CV[i][j]);
    if (A[v] && A[v] != i) {
      cout << 0 << endl;
      return 0;
    }
    A[v] = i;
    done[v] = 1;
    for (int j = 0; j < CV[i].size(); j++) {
      int k = CV[i][j];
      while (!done[k]) {
        A[k] = i;
        done[k] = 1;
        k = lca.parent[0][k];
        if (A[k] && A[k] != i) {
          cout << 0 << endl;
          return 0;
        }
      }
    }
  }
  for (int i = 0; i < N; i++) {
    if (A[i]) {
      dfs2(i, -1);
      cout << dp[i][1] << endl;
      return 0;
    }
  }
  return 0;
}
