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

#include <bits/stdc++.h>

using namespace std;

const int MOD = 1e9 + 7;

// const int MOD=998244353;

long long n;
vector<long long> cols, preds, dp, parents, spec;
vector<vector<int>> tree;
vector<stack<int>> vs, vs2;

void dfs(int x, int p) {
  parents[x] = p;

  vs[cols[p]].push(x);
  if (vs[cols[x]].size())
    preds[x] = vs[cols[x]].top();
  if (!vs2[cols[x]].size())
    spec.push_back(x);
  vs2[cols[x]].push(x);

  for (int i : tree[x]) {
    if (i != p) {
      dfs(i, x);
    }
  }
  vs[cols[p]].pop();
  vs2[cols[x]].pop();
  dp[x] = 1;
  for (int i : tree[x]) {
    if (i != p) {
      dp[x] += dp[i];
    }
  }
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  cin >> n;
  cols.resize(n + 1);
  preds.resize(n + 1);
  tree.resize(n + 1);
  vs.resize(n + 1);
  vs2.resize(n + 1);
  dp.resize(n + 1);
  parents.resize(n + 1);

  for (int i = 1; i <= n; i++)
    cin >> cols[i];
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    tree[x].push_back(y);
    tree[y].push_back(x);
  }
  dfs(1, 0);

  vector<long long> nc = dp;
  for (int i = 1; i <= n; i++) {
    nc[preds[i]] -= dp[i];
  }

  vector<long long> vbase(n + 1, n), ans(n + 1);

  for (int i = 1; i <= n; i++) {
    ans[cols[i]] = (n * (n + 1)) / 2;
  }

  for (int i = 1; i <= n; i++) {
    ans[cols[parents[i]]] -= ((nc[i] + 1) * nc[i]) / 2;
  }

  for (int i : spec) {
    vbase[cols[i]] -= dp[i];
  }

  for (int i = 1; i <= n; i++) {
    if (vbase[i] != n) {
      ans[i] -= (vbase[i] * (vbase[i] + 1)) / 2;
    }
  }

  for (int i = 1; i <= n; i++) {
    cout << ans[i] << '\\n';
  }

  return 0;
}

Error: WA

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

const int MOD = 1e9 + 7;

// const int MOD=998244353;

long long n;
vector<long long> cols, preds, dp, parents, spec;
vector<vector<int>> tree;
vector<stack<int>> vs, vs2;

void dfs(int x, int p) {
  parents[x] = p;

  vs[cols[p]].push(x);
  if (vs[cols[x]].size())
    preds[x] = vs[cols[x]].top();
  if (!vs2[cols[x]].size())
    spec.push_back(x);
  vs2[cols[x]].push(x);

  for (int i : tree[x]) {
    if (i != p) {
      dfs(i, x);
    }
  }
  vs[cols[x]].pop();
  vs2[cols[x]].pop();
  dp[x] = 1;
  for (int i : tree[x]) {
    if (i != p) {
      dp[x] += dp[i];
    }
  }
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  cin >> n;
  cols.resize(n + 1);
  preds.resize(n + 1);
  tree.resize(n + 1);
  vs.resize(n + 1);
  vs2.resize(n + 1);
  dp.resize(n + 1);
  parents.resize(n + 1);

  for (int i = 1; i <= n; i++)
    cin >> cols[i];
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    tree[x].push_back(y);
    tree[y].push_back(x);
  }
  dfs(1, 0);

  vector<long long> nc = dp;
  for (int i = 1; i <= n; i++) {
    nc[preds[i]] -= dp[i];
  }

  vector<long long> vbase(n + 1, n), ans(n + 1);

  for (int i = 1; i <= n; i++) {
    ans[cols[i]] = (n * (n + 1)) / 2;
  }

  for (int i = 1; i <= n; i++) {
    ans[cols[parents[i]]] -= ((nc[i] + 1) * nc[i]) / 2;
  }

  for (int i : spec) {
    vbase[cols[i]] -= dp[i];
  }

  /*
  for(int i=1; i<=n; i++){
      if(vbase[i]!=n){
          ans[i]-=(vbase[i]*(vbase[i]+1))/2;
      }
  }
  */

  for (int i = 1; i <= n; i++) {
    cout << ans[i] << '\\n';
  }

  return 0;
}