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

#include "bits/stdc++.h"
#define hhh cerr << "hhh" << endl
#define see(x) cerr << (#x) << '=' << (x) << endl
using namespace std;
typedef long long ll;
typedef pair<int, int> pr;
inline int read() {
  int x = 0, f = 1;
  char c = getchar();
  while (c != '-' && (c < '0' || c > '9'))
    c = getchar();
  if (c == '-')
    f = -1, c = getchar();
  while (c >= '0' && c <= '9')
    x = x * 10 + c - '0', c = getchar();
  return f * x;
}

const int maxn = 3e5 + 7;
const int inf = 0x3f3f3f3f;
const int mod = 1e9 + 7;

int n;
int head[maxn], to[maxn * 2], nxt[maxn * 2], tot;

inline void add_edge(int u, int v) {
  ++tot;
  to[tot] = v;
  nxt[tot] = head[u];
  head[u] = tot;
  ++tot;
  to[tot] = u;
  nxt[tot] = head[v];
  head[v] = tot;
}

int cnt[maxn], vis[maxn];
int son[maxn], c[maxn], sz[maxn];
ll ans[maxn];

void dfs0(int u, int f) {
  sz[u] = 1;
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f)
      continue;
    dfs0(v, u);
    sz[u] += sz[v];
    if (sz[v] > sz[son[u]])
      son[u] = v;
  }
}

void dfs1(int u, int f, int ff) {
  if (!vis[c[u]])
    cnt[c[u]] += ff ? sz[u] : -sz[u];
  vis[c[u]]++;
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f)
      continue;
    dfs1(v, u, ff);
  }
  vis[c[u]]--;
}

void dfs(int u, int f, int ff) {
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f || v == son[u])
      continue;
    dfs(v, u, 0);
  }
  if (son[u])
    dfs(son[u], u, 1);
  vis[c[u]]++;
  cnt[c[u]] = sz[u];
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f || v == son[u])
      continue;
    dfs1(v, u, 1);
  }
  if (f) {
    ll b = sz[u] - cnt[c[f]];
    ans[c[f]] -= b + b * (b - 1) / 2;
  }
  if (!ff) {
    for (int i = head[u]; i; i = nxt[i]) {
      int v = to[i];
      if (v == f)
        continue;
      dfs1(v, u, 0);
    }
    cnt[c[u]] = 0;
  }
  vis[c[u]]--;
}

int main() {
  n = read();
  for (int i = 1; i <= n; ++i)
    c[i] = read();
  for (int i = 1; i < n; ++i)
    add_edge(read(), read());
  dfs0(1, 0);
  dfs(1, 0, 1);
  for (int i = 1; i <= n; ++i) {
    ll b = n - cnt[i];
    ans[i] -= b + b * (b - 1) / 2;
  }
  ll total = n + ll(n) * (n - 1) / 2;
  for (int i = 1; i <= n; ++i)
    printf("%lld\\n", total + ans[i]);
}

Error: 0

Fixed submission:

#include "bits/stdc++.h"
#define hhh cerr << "hhh" << endl
#define see(x) cerr << (#x) << '=' << (x) << endl
using namespace std;
typedef long long ll;
typedef pair<int, int> pr;
inline int read() {
  int x = 0, f = 1;
  char c = getchar();
  while (c != '-' && (c < '0' || c > '9'))
    c = getchar();
  if (c == '-')
    f = -1, c = getchar();
  while (c >= '0' && c <= '9')
    x = x * 10 + c - '0', c = getchar();
  return f * x;
}

const int maxn = 3e5 + 7;
const int inf = 0x3f3f3f3f;
const int mod = 1e9 + 7;

int n;
int head[maxn], to[maxn * 2], nxt[maxn * 2], tot;

inline void add_edge(int u, int v) {
  ++tot;
  to[tot] = v;
  nxt[tot] = head[u];
  head[u] = tot;
  ++tot;
  to[tot] = u;
  nxt[tot] = head[v];
  head[v] = tot;
}

int cnt[maxn], vis[maxn];
int son[maxn], c[maxn], sz[maxn];
ll ans[maxn];

void dfs0(int u, int f) {
  sz[u] = 1;
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f)
      continue;
    dfs0(v, u);
    sz[u] += sz[v];
    if (sz[v] > sz[son[u]])
      son[u] = v;
  }
}

void dfs1(int u, int f, int ff) {
  if (!vis[c[u]])
    cnt[c[u]] += ff ? sz[u] : -sz[u];
  vis[c[u]]++;
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f)
      continue;
    dfs1(v, u, ff);
  }
  vis[c[u]]--;
}

void dfs(int u, int f, int ff) {
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f || v == son[u])
      continue;
    dfs(v, u, 0);
  }
  int flag = 0;
  if (!vis[c[u]])
    flag = 1;
  if (son[u])
    dfs(son[u], u, 1);
  vis[c[u]]++;
  if (flag)
    cnt[c[u]] = sz[u];
  for (int i = head[u]; i; i = nxt[i]) {
    int v = to[i];
    if (v == f || v == son[u])
      continue;
    dfs1(v, u, 1);
  }
  if (f) {
    ll b = sz[u] - cnt[c[f]];
    ans[c[f]] -= b + b * (b - 1) / 2;
  }
  if (!ff) {
    vis[c[u]]--;
    assert(!vis[c[u]]);
    if (!vis[c[u]])
      cnt[c[u]] = 0;
    for (int i = head[u]; i; i = nxt[i]) {
      int v = to[i];
      if (v == f)
        continue;
      dfs1(v, u, 0);
    }
  }
}

int main() {
  n = read();
  for (int i = 1; i <= n; ++i)
    c[i] = read();
  for (int i = 1; i < n; ++i)
    add_edge(read(), read());
  dfs0(1, 0);
  dfs(1, 0, 1);
  for (int i = 1; i <= n; ++i) {
    ll b = n - cnt[i];
    ans[i] -= b + b * (b - 1) / 2;
  }
  ll total = n + ll(n) * (n - 1) / 2;
  for (int i = 1; i <= n; ++i)
    printf("%lld\\n", total + ans[i]);
}