Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for (int i = 0; i < (n); ++i)
#define repr(i, n) for (int i = (n - 1); i >= 0; --i)

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}

typedef long long ll;
typedef pair<int, int> P;

const int MAX = 200005;
const int INF = 1001001001;
const int MOD = 1000000007;

struct BIT {
  int N, id;
  vector<int> bit;
  vector<pair<int, int>> his;
  BIT(int N, int id = 0) : N(N), id(id), bit(N + 1), his() {
    for (int i = 1; i <= N; ++i)
      add(i, id, false);
  }
  void add(int i, int v, bool flg = true) {
    if (flg)
      his.emplace_back(i, v);
    for (int j = i; j <= N; j += j & -j)
      bit[j] += v;
  }
  int sum(int i) {
    int res = 0;
    for (int j = i; j > 0; j -= j & -j)
      res += bit[j];
    return res;
  }
  int sum(int l, int r) { return sum(r - 1) - sum(l - 1); }
  void clear() {
    for (auto &p : his)
      add(p.first, -p.second);
    his.clear();
  }
};

vector<int> G[MAX];

int dfs(int v, vector<int> &in, vector<int> &out, int p = -1, int k = 1) {
  in[v] = k++;
  for (auto &u : G[v]) {
    if (u == p)
      continue;
    k = dfs(u, in, out, v, k);
  }
  return out[v] = k;
}

ll calc(ll x) { return x * (x + 1) / 2; }

int main() {
  int N;
  cin >> N;

  vector<int> C(N);
  for (auto &i : C)
    cin >> i, --i;

  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  vector<int> in(N), out(N);
  dfs(0, in, out);

  BIT bit = BIT(N, 1);

  vector<vector<int>> V(N);
  rep(i, N) V[C[i]].push_back(i);
  rep(i, N) {
    sort(V[i].begin(), V[i].end(), [&](int a, int b) { return in[a] > in[b]; });
  }

  ll total = calc(N);

  rep(i, N) {
    ll ans = total;
    vector<int> &NV = V[i];
    for (auto &v : NV) {
      int cnt = 1;
      for (auto &u : G[v]) {
        if (in[v] > in[u])
          continue;
        int num = bit.sum(in[u], out[u]);
        ans -= calc(num);
        cnt += num;
      }
      bit.add(in[v], -cnt);
    }
    int num = bit.sum(1, N + 1);
    ans -= calc(num);
    cout << ans << endl;
    bit.clear();
  }
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for (int i = 0; i < (n); ++i)
#define repr(i, n) for (int i = (n - 1); i >= 0; --i)

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}

typedef long long ll;
typedef pair<int, int> P;

const int MAX = 200005;
const int INF = 1001001001;
const int MOD = 1000000007;

struct BIT {
  int N, id;
  vector<int> bit;
  vector<pair<int, int>> his;
  BIT(int N, int id = 0) : N(N), id(id), bit(N + 1), his() {
    for (int i = 1; i <= N; ++i)
      add(i, id, false);
  }
  void add(int i, int v, bool flg = true) {
    if (flg)
      his.emplace_back(i, v);
    for (int j = i; j <= N; j += j & -j)
      bit[j] += v;
  }
  int sum(int i) {
    int res = 0;
    for (int j = i; j > 0; j -= j & -j)
      res += bit[j];
    return res;
  }
  int sum(int l, int r) { return sum(r - 1) - sum(l - 1); }
  void clear() {
    for (auto &p : his)
      add(p.first, -p.second, false);
    his.clear();
  }
};

vector<int> G[MAX];

int dfs(int v, vector<int> &in, vector<int> &out, int p = -1, int k = 1) {
  in[v] = k++;
  for (auto &u : G[v]) {
    if (u == p)
      continue;
    k = dfs(u, in, out, v, k);
  }
  return out[v] = k;
}

ll calc(ll x) { return x * (x + 1) / 2; }

int main() {
  int N;
  cin >> N;

  vector<int> C(N);
  for (auto &i : C)
    cin >> i, --i;

  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  vector<int> in(N), out(N);
  dfs(0, in, out);

  BIT bit = BIT(N, 1);

  vector<vector<int>> V(N);
  rep(i, N) V[C[i]].push_back(i);
  rep(i, N) {
    sort(V[i].begin(), V[i].end(), [&](int a, int b) { return in[a] > in[b]; });
  }

  ll total = calc(N);

  rep(i, N) {
    ll ans = total;
    vector<int> &NV = V[i];
    for (auto &v : NV) {
      int cnt = 1;
      for (auto &u : G[v]) {
        if (in[v] > in[u])
          continue;
        int num = bit.sum(in[u], out[u]);
        ans -= calc(num);
        cnt += num;
      }
      bit.add(in[v], -cnt);
    }
    int num = bit.sum(1, N + 1);
    ans -= calc(num);
    cout << ans << endl;
    bit.clear();
  }
}