Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 1000000007
#define MOD2 998244353
#define INF (1 << 29)
#define LINF (1LL << 60)
#define EPS (1e-10)
#define PI 3.1415926535897932384626433832795028
typedef long long Int;
typedef pair<Int, Int> P;
typedef long double Real;
typedef complex<Real> CP;

Int n;
vector<Int> edge[220000];
Int c[220000];
Int ans[220000];
int last_c[220000];
Int tree_size[220000];
Int parent[220000];
Int no_parent_cnt[220000];
vector<int> children[22000];

Int cnt(Int a) { return a * (a - 1) / 2 + a; }

void dfs(int x, int last = -1) {
  parent[x] = last_c[c[x]];
  last_c[c[x]] = x;
  if (parent[x] != -1) {
    children[parent[x]].push_back(x);
  }
  tree_size[x] = 1;
  for (auto to : edge[x]) {
    if (to == last)
      continue;
    dfs(to, x);
    tree_size[x] += tree_size[to];
    Int tmp = tree_size[to];
    for (auto child : children[x]) {
      tmp -= tree_size[child];
    }
    children[x].clear();
    ans[c[x]] -= cnt(tmp);
  }
  if (parent[x] == -1) {
    no_parent_cnt[c[x]] -= tree_size[x];
  }
  last_c[c[x]] = parent[x];
}

int main() {
  cin >> n;
  fill(last_c, last_c + 220000, -1);
  fill(no_parent_cnt, no_parent_cnt + 220000, n);
  fill(ans, ans + n, cnt(n));
  for (int i = 0; i < n; i++) {
    cin >> c[i];
    c[i]--;
  }
  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    edge[u].push_back(v);
    edge[v].push_back(u);
  }
  dfs(0);
  for (int i = 0; i < n; i++) {
    ans[i] -= cnt(no_parent_cnt[i]);
    cout << ans[i] << endl;
  }
  return 0;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 1000000007
#define MOD2 998244353
#define INF (1 << 29)
#define LINF (1LL << 60)
#define EPS (1e-10)
#define PI 3.1415926535897932384626433832795028
typedef long long Int;
typedef pair<Int, Int> P;
typedef long double Real;
typedef complex<Real> CP;

Int n;
vector<Int> edge[220000];
Int c[220000];
Int ans[220000];
int last_c[220000];
Int tree_size[220000];
Int parent[220000];
Int no_parent_cnt[220000];
vector<int> children[220000];

Int cnt(Int a) { return a * (a - 1) / 2 + a; }

void dfs(int x, int last = -1) {
  parent[x] = last_c[c[x]];
  last_c[c[x]] = x;
  if (parent[x] != -1) {
    children[parent[x]].push_back(x);
  }
  tree_size[x] = 1;
  for (auto to : edge[x]) {
    if (to == last)
      continue;
    dfs(to, x);
    tree_size[x] += tree_size[to];
    Int tmp = tree_size[to];
    for (auto child : children[x]) {
      tmp -= tree_size[child];
    }
    children[x].clear();
    ans[c[x]] -= cnt(tmp);
  }
  if (parent[x] == -1) {
    no_parent_cnt[c[x]] -= tree_size[x];
  }
  last_c[c[x]] = parent[x];
}

int main() {
  cin >> n;
  fill(last_c, last_c + 220000, -1);
  fill(no_parent_cnt, no_parent_cnt + 220000, n);
  fill(ans, ans + n, cnt(n));
  for (int i = 0; i < n; i++) {
    cin >> c[i];
    c[i]--;
  }
  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    edge[u].push_back(v);
    edge[v].push_back(u);
  }
  dfs(0);
  for (int i = 0; i < n; i++) {
    ans[i] -= cnt(no_parent_cnt[i]);
    cout << ans[i] << endl;
  }
  return 0;
}