Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define int long long
struct ss {
  int node, nxt;
} e[600005];
inline int read() {
  int re = 0, k = 1;
  char ch = getchar();
  while (ch < '0' || ch > '9') {
    if (ch == '-')
      k = -1;
    ch = getchar();
  }
  while (ch >= '0' && ch <= '9') {
    re = (re << 1) + (re << 3) + ch - 48;
    ch = getchar();
  }
  return re * k;
}
int u, v, n, col[100005], f[100005], del[100005], s[100005], sz[100005],
    head[100005], tot, g[100005];
void add(int u, int v) {
  e[++tot].nxt = head[u];
  e[tot].node = v;
  head[u] = tot;
}
void getsz(int x, int fa) {
  for (int i = head[x]; i; i = e[i].nxt) {
    int t = e[i].node;
    if (t == fa)
      continue;
    getsz(t, x);
    sz[x] += sz[t];
  }
}
void dfs(int x, int fa) {
  int lst = s[col[fa]];
  s[col[fa]] = x;
  if (s[col[x]])
    del[s[col[x]]] += sz[x];
  else
    g[col[x]] += sz[x];
  for (int i = head[x]; i; i = e[i].nxt) {
    int t = e[i].node;
    if (t == fa)
      continue;
    dfs(t, x);
  }
  int now = sz[x] - del[x];
  // cerr<<col[fa]<<" "<<now<<" "<<sz[x]<<endl;
  f[col[fa]] += now * (now + 1) >> 1;
  s[col[fa]] = lst;
}
signed main() {
  n = read();
  for (int i = 1; i <= n; i++) {
    col[i] = read();
    sz[i] = 1;
  }
  for (int i = 1; i <= n - 1; i++) {
    u = read();
    v = read();
    add(u, v);
    add(v, u);
  }
  getsz(1, 0);
  dfs(1, 0);
  for (int i = 1; i <= n; i++)
    g[i] = (n - g[i]) * (n - g[i] + 1) >> 1;
  int ans = n * (n + 1) >> 1;
  for (int i = 1; i <= n; i++) {
    printf("%lld\n", ans - f[i] - g[i]);
  }
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define int long long
struct ss {
  int node, nxt;
} e[600005];
inline int read() {
  int re = 0, k = 1;
  char ch = getchar();
  while (ch < '0' || ch > '9') {
    if (ch == '-')
      k = -1;
    ch = getchar();
  }
  while (ch >= '0' && ch <= '9') {
    re = (re << 1) + (re << 3) + ch - 48;
    ch = getchar();
  }
  return re * k;
}
int u, v, n, col[200005], f[200005], del[200005], s[200005], sz[200005],
    head[200005], tot, g[200005];
void add(int u, int v) {
  e[++tot].nxt = head[u];
  e[tot].node = v;
  head[u] = tot;
}
void getsz(int x, int fa) {
  for (int i = head[x]; i; i = e[i].nxt) {
    int t = e[i].node;
    if (t == fa)
      continue;
    getsz(t, x);
    sz[x] += sz[t];
  }
}
void dfs(int x, int fa) {
  int lst = s[col[fa]];
  s[col[fa]] = x;
  if (s[col[x]])
    del[s[col[x]]] += sz[x];
  else
    g[col[x]] += sz[x];
  for (int i = head[x]; i; i = e[i].nxt) {
    int t = e[i].node;
    if (t == fa)
      continue;
    dfs(t, x);
  }
  int now = sz[x] - del[x];
  // cerr<<col[fa]<<" "<<now<<" "<<sz[x]<<endl;
  f[col[fa]] += now * (now + 1) >> 1;
  s[col[fa]] = lst;
}
signed main() {
  n = read();
  for (int i = 1; i <= n; i++) {
    col[i] = read();
    sz[i] = 1;
  }
  for (int i = 1; i <= n - 1; i++) {
    u = read();
    v = read();
    add(u, v);
    add(v, u);
  }
  getsz(1, 0);
  dfs(1, 0);
  for (int i = 1; i <= n; i++)
    g[i] = (n - g[i]) * (n - g[i] + 1) >> 1;
  int ans = n * (n + 1) >> 1;
  for (int i = 1; i <= n; i++) {
    printf("%lld\n", ans - f[i] - g[i]);
  }
}