#include <bits/stdc++.h>
using namespace std;
#pragma GCC optimize("O3")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("sse2")
inline void opting() {
  ios_base::sync_with_stdio(false);
  cout.tie(0);
  cin.tie(0);
}
long long inp() {
  long long x;
  cin >> x;
  return x;
}
const long long maxn = 3e5 + 69, delta = 998244353;
vector<long long> adj[maxn];
void add_edge(long long u, long long v) {
  adj[v].push_back(u);
  adj[u].push_back(v);
}
pair<long long, long long> min_maxing[maxn];
long long st[maxn], en[maxn];
long long color[maxn];
long long dp1[maxn], dp2[maxn];
bool poss = true;
long long inverse(long long A) {
  A %= delta;
  long long b = delta - 2;
  long long z = A;
  long long ans = 1;
  while (b) {
    if (b & 1) ans = (ans * z) % delta;
    b /= 2;
    z = (z * z) % delta;
  }
  return ans;
}
void color_nodes(long long v, long long p) {
  for (auto u : adj[v])
    if (u != p) {
      color_nodes(u, v);
      if (color[u] != 0 and color[u] != color[v] and
          !(st[u] <= min_maxing[color[u]].first and
            min_maxing[color[u]].second <= en[u]))
        if (color[v] != 0)
          poss = false;
        else
          color[v] = color[u];
    }
}
void wi_color(long long v, long long p) {
  dp2[v] = 0;
  long long X = 1;
  for (auto u : adj[v])
    if (u != p)
      if (color[u] == color[v])
        X = (X * dp1[u]) % delta;
      else
        X = (X * (dp1[u] + dp2[u])) % delta;
  dp1[v] = X;
}
void no_color(long long v, long long p) {
  long long X = 1;
  for (auto u : adj[v])
    if (u != p) X = (X * (dp1[u] + dp2[u])) % delta;
  dp2[v] = X;
  for (auto u : adj[v])
    if (u != p)
      dp1[v] =
          (dp1[v] + dp1[u] * ((X * inverse(dp1[u] + dp2[u])) % delta)) % delta;
}
void DP_calc(long long v, long long p) {
  for (auto u : adj[v])
    if (u != p) DP_calc(u, v);
  if (color[v] != 0)
    wi_color(v, p);
  else
    no_color(v, p);
  long long i = v;
}
long long n;
void input() {
  n = inp();
  long long k = inp();
  fill(min_maxing, min_maxing + k + 5, make_pair(maxn, -1));
  for (long long i = 0; i < n; i++) cin >> color[i];
  for (long long i = 1; i < n; i++) add_edge(inp() - 1, inp() - 1);
}
long long TIME = 0;
void dfs1(long long v, long long p) {
  st[v] = TIME++;
  for (auto u : adj[v])
    if (u != p) dfs1(u, v);
  en[v] = TIME;
  min_maxing[color[v]].first = min(min_maxing[color[v]].first, st[v]);
  min_maxing[color[v]].second = max(min_maxing[color[v]].second, en[v]);
}
void solve() {
  dfs1(0, 0);
  color_nodes(0, 0);
  DP_calc(0, 0);
  cout << poss * dp1[0] << endl;
}
void debug() {
  for (long long i = 0; i < n; i++)
    cout << i + 1 << ' ' << color[i] << ' ' << dp1[i] << ' ' << dp2[i] << endl;
}
int main() {
  opting();
  input();
  solve();
}
