#include <bits/stdc++.h>
using namespace std;
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  long long n, k;
  cin >> n >> k;
  vector<long long> a((size_t)n);
  vector<vector<long long>> a_reversed((size_t)k + 1);
  for (long long u = 0; u < n; u++) {
    cin >> a[u];
    a_reversed[a[u]].push_back(u);
  }
  vector<vector<long long>> adj((size_t)n);
  for (long long i = 0; i < n - 1; i++) {
    long long u, v;
    cin >> u >> v;
    u--, v--;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }
  vector<long long> vertex_parents((size_t)n);
  vector<long long> vertex_in((size_t)n);
  vector<long long> vertex_out((size_t)n);
  vector<vector<long long>> subtree_colors((size_t)k + 1);
  long long time = 0;
  function<void(long long, long long)> visit_initial = [&](long long u,
                                                           long long v) {
    vertex_parents[v] = u;
    vertex_in[v] = time;
    subtree_colors[a[v]].push_back(time);
    time++;
    for (long long w : adj[v]) {
      if (w != u) {
        visit_initial(v, w);
      }
    }
    vertex_out[v] = time - 1;
  };
  visit_initial(-1, 0);
  function<long long(long long, long long)> subtree_size = [&](long long u,
                                                               long long v) {
    if (u != -1 && v == vertex_parents[u]) {
      return n - subtree_size(v, u);
    } else {
      return vertex_out[v] - vertex_in[v] + 1;
    }
  };
  function<long long(long long, long long, long long)> subtree_color =
      [&](long long u, long long v, long long c) {
        if (u != -1 && v == vertex_parents[u]) {
          return (long long)a_reversed[c].size() - subtree_color(v, u, c);
        } else {
          return (
              long long)(upper_bound(subtree_colors[c].begin(),
                                     subtree_colors[c].end(), vertex_out[v]) -
                         lower_bound(subtree_colors[c].begin(),
                                     subtree_colors[c].end(), vertex_in[v]));
        }
      };
  function<void(long long, long long, long long)> visit_check =
      [&](long long u, long long v, long long c) {
        if (a[v] == 0) {
          a[v] = c;
        } else if (a[v] != c) {
          cout << 0 << endl;
          exit(0);
        }
        for (long long w : adj[v]) {
          if (w != u) {
            if (subtree_color(v, w, c) > 0 ||
                subtree_color(v, w, 0) == subtree_size(v, w)) {
              visit_check(v, w, c);
            }
          }
        }
      };
  for (long long c = 1; c <= k; c++) {
    visit_check(-1, a_reversed[c][0], c);
  }
  set<long long> remaining_vertexes;
  for (long long u = 0; u < n; u++) {
    if (a[u] == 0) {
      for (long long v : adj[u]) {
        if (a[v] != 0) {
          remaining_vertexes.insert(u);
        }
      }
    }
  }
  function<long long(long long, long long)> power = [&](long long x,
                                                        long long y) {
    if (y == 0) {
      return (long long)1;
    } else {
      long long z = power(x, y / 2);
      return y % 2 == 0 ? (z * z) % 998244353
                        : (((z * z) % 998244353) * x) % 998244353;
    }
  };
  auto div = [&](long long x) { return power(x, 998244353 - 2); };
  function<pair<long long, long long>(long long, long long)> visit_count =
      [&](long long u, long long v) {
        remaining_vertexes.erase(v);
        if (a[v] != 0) {
          return pair<long long, long long>{1, 0};
        } else {
          pair<long long, long long> x{0, 1};
          set<long long> s;
          vector<pair<long long, long long>> t;
          for (long long w : adj[v]) {
            if (w != u) {
              if (a[v] == 0 || s.insert(a[v]).second) {
                pair<long long, long long> y = visit_count(v, w);
                x.second *= y.first + y.second;
                x.second %= 998244353;
                t.push_back(y);
              }
            }
          }
          for (pair<long long, long long> y : t) {
            x.first +=
                (x.second * div(y.first + y.second)) % 998244353 * y.first;
            x.first %= 998244353;
          }
          return x;
        }
      };
  long long o = 1;
  while (!remaining_vertexes.empty()) {
    long long u = *remaining_vertexes.begin();
    remaining_vertexes.erase(remaining_vertexes.begin());
    o *= visit_count(-1, u).first;
    o %= 998244353;
  }
  cout << o << endl;
  return 0;
}
