Problem p02710 (C++)

Buggy submission (status: Runtime Error):
// <head>
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int INF = 1002003004;
const ll LINF = 1002003004005006007ll;
struct preprocess {
  preprocess() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);
  }
} ____;
// </head>

// <library>
template <typename T> struct BIT {
  int n;
  vector<T> d;
  BIT(int n = 0) : n(n), d(n + 1) {}
  void add(int i, T x = 1) {
    for (i++; i <= n; i += i & -i) {
      d[i] += x;
    }
  }
  T sum(int i) {
    T x = 0;
    for (i++; i; i -= i & -i) {
      x += d[i];
    }
    return x;
  }
  T sum(int l, int r) { return sum(r - 1) - sum(l - 1); }
};
// </library>
const int MX = 10;
int in[MX], out[MX], k;
vector<int> to[MX];
void dfs(int v, int p = -1) {
  in[v] = k++;
  for (int u : to[v]) {
    if (u == p)
      continue;
    dfs(u, v);
  }
  out[v] = k;
}
vector<ll> vs[MX];
ll f(ll n) { return n * (n + 1) / 2; }
int main() {
  int n;
  cin >> n;
  vector<ll> c(n);

  for (int i = 0; i < n; i++) {
    cin >> c[i];
    vs[c[i] - 1].emplace_back(i);
  }
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    to[a].emplace_back(b);
    to[b].emplace_back(a);
  }
  dfs(0);
  ll tot = f(n);
  BIT<int> d(n);
  for (int i = 0; i < n; i++) {
    d.add(i, 1);
  }
  for (int i = 0; i < n; i++) {
    ll ans = tot;
    vector<ll> nvs = vs[i];
    sort(nvs.begin(), nvs.end(), [&](int a, int b) { return in[a] > in[b]; });
    vector<pair<ll, ll>> history;
    for (auto j : nvs) {
      ll cnt = 1;
      for (auto k : to[j]) {
        if (in[j] > in[k])
          continue;
        ll num = d.sum(in[k], out[k]);
        ans -= f(num);
        cnt += num;
      }
      d.add(in[j], -cnt);
      history.emplace_back(make_pair(in[j], cnt));
    }
    ans -= f(d.sum(0, n));
    cout << ans << '\n';
    for (auto p : history) {
      d.add(p.first, p.second);
    }
  }
}

Fixed submission:
// <head>
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int INF = 1002003004;
const ll LINF = 1002003004005006007ll;
struct preprocess {
  preprocess() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);
  }
} ____;
// </head>

// <library>
template <typename T> struct BIT {
  int n;
  vector<T> d;
  BIT(int n = 0) : n(n), d(n + 1) {}
  void add(int i, T x = 1) {
    for (i++; i <= n; i += i & -i) {
      d[i] += x;
    }
  }
  T sum(int i) {
    T x = 0;
    for (i++; i; i -= i & -i) {
      x += d[i];
    }
    return x;
  }
  T sum(int l, int r) { return sum(r - 1) - sum(l - 1); }
};
// </library>
const int MX = 200005;
int in[MX], out[MX], k;
vector<int> to[MX];
void dfs(int v, int p = -1) {
  in[v] = k++;
  for (int u : to[v]) {
    if (u == p)
      continue;
    dfs(u, v);
  }
  out[v] = k;
}
vector<ll> vs[MX];
ll f(ll n) { return n * (n + 1) / 2; }
int main() {
  int n;
  cin >> n;
  vector<ll> c(n);

  for (int i = 0; i < n; i++) {
    cin >> c[i];
    vs[c[i] - 1].emplace_back(i);
  }
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    to[a].emplace_back(b);
    to[b].emplace_back(a);
  }
  dfs(0);
  ll tot = f(n);
  BIT<int> d(n);
  for (int i = 0; i < n; i++) {
    d.add(i, 1);
  }
  for (int i = 0; i < n; i++) {
    ll ans = tot;
    vector<ll> nvs = vs[i];
    sort(nvs.begin(), nvs.end(), [&](int a, int b) { return in[a] > in[b]; });
    vector<pair<ll, ll>> history;
    for (auto j : nvs) {
      ll cnt = 1;
      for (auto k : to[j]) {
        if (in[j] > in[k])
          continue;
        ll num = d.sum(in[k], out[k]);
        ans -= f(num);
        cnt += num;
      }
      d.add(in[j], -cnt);
      history.emplace_back(make_pair(in[j], cnt));
    }
    ans -= f(d.sum(0, n));
    cout << ans << '\n';
    for (auto p : history) {
      d.add(p.first, p.second);
    }
  }
}