Submission #50218488

---

Source Code Expand

#include <iostream>
#include <vector>
#include <unordered_map>
#include <atcoder/modint>
using namespace std;
using mint = atcoder::static_modint<998244353>;

int main(){
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
  int n;
  cin >> n;
  vector<int> a(n);
  for(int i = 0; i < n; i++) cin >> a[i];
  vector<vector<int>> g(n);
  // グラフの作成
  for(int i = 0; i < n - 1; i++){
    int x, y;
    cin >> x >> y;
    x--, y--;
    g[x].push_back(y);
    g[y].push_back(x);
  }

  vector<unordered_map<int, mint>> D(n);
  mint ans = 0;
  // mpにキーがkである要素がある場合その値を, ない場合0を返す
  auto get = [](unordered_map<int, mint> &mp, int k) -> mint {
    auto itr = mp.find(k);
    return itr == mp.end() ? 0 : itr->second;
  };
  // mpにキーがkである要素がある場合xを足す, ない場合xを追加
  auto add = [](unordered_map<int, mint> &mp, int k, mint x) -> void {
    auto itr = mp.find(k);
    if(itr == mp.end()) mp.emplace(k, x);
    else itr->second += x;
  };
  // 木dp
  auto f = [&](auto &&f, int v, int p) -> void {
    for(int c : g[v]){
      if(c == p) continue;
      f(f, c, v);
      ans += get(D[c], a[v]); // (2)
      
      // |D[v]| >= |D[c]|にする(マージテク)
      if(D[v].size() < D[c].size()) D[v].swap(D[c]);
      
      for(auto [color, cnt] : D[c]){
        mint x = get(D[v], color);
        mint y = cnt * x;
        ans += y; // (3)
        add(D[v], color, y + cnt);
      }
    }
    ans++; // (1)
    add(D[v], a[v], 1);
  };
  f(f, 0, -1);
  cout << ans.val() << '\n';
}

Submission Info

Judge Result