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

#include <bits/stdc++.h>
using namespace std;
#define all(V) V.begin(), V.end()
#define ll long long
#define MOD 1000000007

vector<int> G[200010];
stack<ll> col[200010];
int C[200010];
ll ans[200010];

ll seg[524288];
ll app[200010];

void raq(int s = 0, int t = 262143, ll v = 1, int k = 1, int l = 0,
         int r = 262143) {
  if (t < l || r < s)
    return;
  if (s <= l && r <= t) {
    seg[k] += v;
    return;
  }
  raq(s, t, v, k * 2, l, (l + r) / 2);
  raq(s, t, v, k * 2 + 1, (l + r) / 2 + 1, r);
}

ll gets(int x) {
  x += 262144;
  ll ans = 0;
  while (x) {
    ans += seg[x];
    x /= 2;
  }
  return ans;
}

void pini(int x, ll v) { raq(x, x, v - gets(x)); }

ll alc = 0;

void euler(int c, int p) {

  for (int i : G[c]) {
    if (i == p)
      continue;
    col[C[c]].top() = gets(C[c]);
    col[C[c]].push(0);
    pini(C[c], 0);
    raq();
    raq(C[i], C[i], -1);
    euler(i, c);
  }

  if (p == 0)
    return;
  ll gc = gets(C[p]);
  ans[C[p]] += gc * (gc - 1LL) / 2LL;
  col[C[p]].pop();
  pini(C[p], col[C[p]].top());
  return;
}

int main() {
  for (int i = 0; i < 200010; i++)
    ans[i] = app[i] = C[i] = 0;
  for (int i = 0; i < 524288; i++)
    seg[i] = 0;
  ll N, a, b;
  cin >> N;
  for (int i = 1; i <= N; i++) {
    cin >> C[i];
    app[C[i]]++;
  }
  for (int i = 1; i < N; i++) {
    cin >> a >> b;
    G[a].push_back(b);
    G[b].push_back(a);
  }
  int st = 0;
  for (int i = 1; i <= N; i++) {
    if (G[i].size() == 1)
      st = i;
  }
  for (int i = 1; i <= N; i++) {
    if (i == C[st])
      col[i].push(0);
    col[i].push(1);
  }
  raq();
  raq(C[st], C[st], -1);
  euler(st, 0);
  for (int i = 1; i <= N; i++) {
    ll gc = gets(i);
    ans[i] += gc * (gc - 1LL) / 2LL;
  }
  for (int i = 1; i <= N; i++) {
    if (app[i] == 0)
      cout << 0 << endl;
    else
      cout << N * (N - 1) / 2 - ans[i] + app[i] << endl;
  }
}


Error: -6
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc


Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define all(V) V.begin(), V.end()
#define ll long long
#define MOD 1000000007

vector<int> G[200010];
stack<ll> col[200010];
int C[200010];
ll ans[200010];

ll seg[524288];
ll app[200010];

void raq(int s = 0, int t = 262143, ll v = 1, int k = 1, int l = 0,
         int r = 262143) {
  if (t < l || r < s)
    return;
  if (s <= l && r <= t) {
    seg[k] += v;
    return;
  }
  raq(s, t, v, k * 2, l, (l + r) / 2);
  raq(s, t, v, k * 2 + 1, (l + r) / 2 + 1, r);
}

ll gets(int x) {
  x += 262144;
  ll ans = 0;
  while (x) {
    ans += seg[x];
    x /= 2;
  }
  return ans;
}

void pini(int x, ll v) { raq(x, x, v - gets(x)); }

ll alc = 0;

void euler(int c, int p) {

  for (int i : G[c]) {
    if (i == p)
      continue;
    col[C[c]].top() = gets(C[c]);
    col[C[c]].push(0);
    pini(C[c], 0);
    raq();
    raq(C[i], C[i], -1);
    euler(i, c);
  }

  if (p == 0)
    return;
  ll gc = gets(C[p]);
  ans[C[p]] += gc * (gc - 1LL) / 2LL;
  col[C[p]].pop();
  pini(C[p], col[C[p]].top());
  return;
}

int main() {
  for (int i = 0; i < 200010; i++)
    ans[i] = app[i] = C[i] = 0;
  for (int i = 0; i < 524288; i++)
    seg[i] = 0;
  ll N, a, b;
  cin >> N;
  for (int i = 1; i <= N; i++) {
    cin >> C[i];
    app[C[i]]++;
  }
  for (int i = 1; i < N; i++) {
    cin >> a >> b;
    G[a].push_back(b);
    G[b].push_back(a);
  }
  int st = 0;
  for (int i = 1; i <= N; i++) {
    if (G[i].size() == 1)
      st = i;
  }
  if (st == 0)
    exit(1);
  for (int i = 1; i <= N; i++) {
    if (i == C[st])
      col[i].push(0);
    col[i].push(1);
  }
  raq();
  raq(st, st, -1);
  euler(st, 0);
  for (int i = 1; i <= N; i++) {
    ll gc = gets(i);
    ans[i] += gc * (gc - 1LL) / 2LL;
  }
  for (int i = 1; i <= N; i++) {
    if (app[i] == 0)
      cout << 0 << endl;
    else
      cout << N * (N - 1) / 2 - ans[i] + app[i] << endl;
  }
}