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

#include <algorithm>
#include <stdio.h>
#include <utility>
#include <vector>
using namespace std;
typedef long long ll;

vector<ll> vec[200000], cntt[200000];
vector<pair<ll, ll>> col[200000];
ll n, p[200000], c[200000], r[200000], cnt[200000];
pair<ll, ll> tb[200000], cb[200000];

ll comb(ll a) { return a * (a + 1) / 2; }

void dfs(ll a) {
  ll i;
  pair<ll, ll> tmp = cb[c[a]];
  tb[a] = tmp, cb[c[a]].first = a;
  for (i = 0; i < vec[a].size(); ++i)
    if (p[a] != vec[a][i])
      cb[c[a]].second = i, p[vec[a][i]] = a, r[vec[a][i]] = r[a] - 1,
      dfs(vec[a][i]);
  if (a)
    cnt[p[a]] += cnt[a];
  cb[c[a]] = tmp;
}

int main(void) {
  ll i, j, k, a, b, ans, tmp, now;
  scanf("%lld", &n);
  for (i = 0; i < n; ++i)
    cb[i] = make_pair(-1, 0), cnt[i] = 1;
  p[0] = -1;
  for (i = 0; i < n; ++i)
    scanf("%lld", &c[i]), c[i]--;
  for (i = 1; i < n; ++i) {
    scanf("%lld%lld", &a, &b);
    vec[--a].push_back(--b);
    vec[b].push_back(a);
  }
  r[0] = 0;
  dfs(0);
  for (i = 0; i < n; ++i)
    col[c[i]].push_back(make_pair(r[i], i));
  for (i = 0; i < n; ++i)
    for (j = 0; j < vec[i].size(); ++j) {
      if (p[i] == vec[i][j])
        cntt[i].push_back(0);
      else
        cntt[i].push_back(cnt[vec[i][j]]);
    }
  for (i = 0; i < n; ++i)
    sort(col[i].begin(), col[i].end());
  for (i = 0; i < n; ++i) {
    ans = comb(n), tmp = n;
    for (j = 0; j < col[i].size(); ++j) {
      now = col[i][j].second;
      for (k = 0; k < cntt[now].size(); ++k)
        ans -= comb(cntt[now][k]);
      if (tb[now].first != -1)
        cntt[tb[now].first][tb[now].second] -= cnt[now];
      else
        tmp -= cnt[now];
    }
    ans -= comb(tmp);
    printf("%lld\\n", ans);
  }
  return 0;
}

Error: 0

Fixed submission:

#include <algorithm>
#include <stdio.h>
#include <utility>
#include <vector>
using namespace std;
typedef long long ll;

vector<ll> vec[200000], cntt[200000];
vector<pair<ll, ll>> col[200000];
ll n, p[200000], c[200000], r[200000], cnt[200000];
pair<ll, ll> tb[200000], cb[200000];

ll comb(ll a) { return a * (a + 1) / 2; }

void dfs(ll a) {
  ll i;
  pair<ll, ll> tmp = cb[c[a]];
  tb[a] = tmp, cb[c[a]].first = a;
  for (i = 0; i < vec[a].size(); ++i)
    if (p[a] != vec[a][i])
      cb[c[a]].second = i, p[vec[a][i]] = a, r[vec[a][i]] = r[a] - 1,
      dfs(vec[a][i]);
  cnt[p[a]] += cnt[a];
  cb[c[a]] = tmp;
}

int main(void) {
  ll i, j, k, a, b, ans, tmp, sum, now;
  scanf("%lld", &n);
  for (i = 0; i < n; ++i)
    tb[i] = cb[i] = make_pair(-1, 0), cnt[i] = 1;
  p[0] = -1;
  for (i = 0; i < n; ++i)
    scanf("%lld", &c[i]), c[i]--;
  for (i = 1; i < n; ++i) {
    scanf("%lld%lld", &a, &b);
    vec[--a].push_back(--b);
    vec[b].push_back(a);
  }
  r[0] = 0;
  dfs(0);
  for (i = 0; i < n; ++i)
    col[c[i]].push_back(make_pair(r[i], i));
  for (i = 0; i < n; ++i)
    for (j = 0; j < vec[i].size(); ++j) {
      if (p[i] == vec[i][j])
        cntt[i].push_back(0);
      else
        cntt[i].push_back(cnt[vec[i][j]]);
    }
  for (i = 0; i < n; ++i)
    sort(col[i].begin(), col[i].end());
  for (i = 0; i < n; ++i) {
    ans = comb(n), tmp = n;
    for (j = 0; j < col[i].size(); ++j) {
      now = col[i][j].second;
      for (k = 0; k < cntt[now].size(); ++k)
        ans -= comb(cntt[now][k]);
      if (tb[now].first != -1)
        cntt[tb[now].first][tb[now].second] -= cnt[now];
      else
        tmp -= cnt[now];
    }
    ans -= comb(tmp);
    printf("%lld\\n", ans);
  }
  return 0;
}