Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
ll dfs(ll x, vector<vector<ll>> &g, vector<bool> &see, vector<ll> &cl,
       vector<ll> &ans, vector<ll> &c, vector<stack<ll>> &cn, vector<ll> &par) {
  ll chi = 0;
  // 点xに対する処理
  see[x] = 1;
  ll p = -1;
  if (!cn[c[x]].empty())
    p = cn[c[x]].top();
  cn[c[x]].push(x);
  for (auto r : g[x]) {
    if (see[r])
      continue;
    // 進む辺に対する処理（これは木構造でオイラーツアーするために使う）
    cl[x] = 0;
    ll chit = dfs(r, g, see, cl, ans, c, cn, par) + 1;
    ans[c[x]] -= (chit + cl[x]) * (chit + cl[x] + 1) / 2;
    // 戻る辺に対する処理
    chi += chit;
  }
  // 点xの親に戻る時の処理
  if (p == -1) {
    par[c[x]] -= chi + 1;
  } else {
    cl[p] -= chi + 1;
  }
  cn[c[x]].pop();
  return chi;
}

int main() {
  ll n;
  cin >> n;
  vector<ll> ans(n, n * (n + 1) / 2);
  vector<ll> cl(n, 0);
  vector<ll> par(n, n);
  vector<stack<ll>> cn(n);
  vector<bool> see(n, 0);
  vector<ll> c(n, 0);
  rep(i, n) {
    ll cc;
    cin >> cc;
    c[i] = cc - 1;
  }
  vector<vector<ll>> g(n);
  rep(i, n - 1) {
    ll a, b;
    cin >> a >> b;
    g[a - 1].push_back(b - 1);
    g[b - 1].push_back(a - 1);
  }
  dfs(0, g, see, cl, ans, c, cn, par);
  rep(i, n) cout << ans[i] - par[i] * (par[i] + 1) / 2 << endl;
  return 0;
}

Error: 0

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
vector<vector<ll>> num(100010);
vector<queue<ll>> cl(100010);
vector<ll> c(100010);
vector<ll> itr(100010);

void dfs(ll x, vector<vector<ll>> &g, vector<bool> &see, vector<ll> &numc) {
  // 点xに対する処理
  if (!cl[c[x]].empty()) {
    ll p = cl[c[x]].front();
    num[p][itr[p]] -= numc[x] + 1;
  }
  cl[c[x]].push(x);
  see[x] = 1;
  for (auto p : g[x]) {
    if (see[p])
      continue;
    // 進む辺に対する処理（これは木構造でオイラーツアーするために使う）
    dfs(p, g, see, numc);
    itr[x]++;
    // 戻る辺に対する処理
  }
  // 点xの親に戻る時の処理
  cl[c[x]].pop();
}

int main() {
  ll n;
  cin >> n;
  vector<ll> ans(n, 0);
  rep(i, n) { cin >> c[i]; }
  vector<vector<ll>> g(n);
  rep(i, n - 1) {
    ll a, b;
    cin >> a >> b;
    g[a - 1].push_back(b - 1);
    g[b - 1].push_back(a - 1);
  }

  vector<vector<ll>> child(n);
  vector<ll> numc(n);
  vector<ll> par(n);
  queue<ll> q;
  queue<ll> qq;
  ll root = 0;
  q.push(root);
  par[root] = -1;
  // 親の定義と葉を見つける
  while (!q.empty()) {
    ll p = q.front();
    q.pop();
    if (g[p].size() == 1 && p != root) {
      qq.push(p);
    }
    for (auto r : g[p]) {
      if (r == par[p])
        continue;
      par[r] = p;
      q.push(r);
    }
  }
  while (!qq.empty()) {
    ll p = qq.front();
    qq.pop();
    child[par[p]].push_back(p); // 親に値を追加する
    // 根の時は例外処理
    if (child[par[p]].size() == g[par[p]].size()) {
      // child[par[p]]に子の値が入ってる

      ll tm = 0;
      for (auto r : child[par[p]]) {
        tm += 1 + numc[r];
        num[par[p]].push_back(1 + numc[r]);
      }
      numc[par[p]] = tm;
    }
    // 全ての葉が親に追加し終わった
    if (child[par[p]].size() == g[par[p]].size() - 1 && par[p] != root) {
      // child[par[p]]に子の値が入ってる

      ll tm = 0;
      for (auto r : child[par[p]]) {
        tm += 1 + numc[r];
        num[par[p]].push_back(1 + numc[r]);
      }
      numc[par[p]] = tm;

      qq.push(par[p]);
    }
  }
  rep(i, n) { num[i].push_back(n - 1 - numc[i]); }

  vector<bool> see(n);
  dfs(0, g, see, numc);

  rep(i, n) {
    ll sum = 0;
    for (auto r : num[i]) {
      sum += r;
    }
    for (auto r : num[i]) {
      ans[c[i] - 1] += r * (sum - r);
    }
  }
  rep(i, n) ans[i] /= 2;
  rep(i, n) {
    ll sum = 0;
    for (auto r : num[i]) {
      sum += r;
    }
    ans[c[i] - 1] += sum + 1;
  }

  for (auto an : ans)
    cout << an << endl;
  return 0;
}