Problem p02710 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

#ifdef DEBUG
void debug() { cerr << "\\n"; }
template <class T> void debug(const T &x) { cerr << x << "\\n"; }
template <class T, class... Args> void debug(const T &x, const Args &...args) {
  cerr << x << " ";
  debug(args...);
}
template <class T> void debugVector(const vector<T> &v) {
  for (const T &x : v) {
    cerr << x << " ";
  }
  cerr << "\\n";
}
#else
template <class T, class... Args> void debug(const T &x, const Args &...args) {}
template <class T> void debugVector(const vector<T> &v) {}
#endif

using ll = long long;
#define ALL(v) (v).begin(), (v).end()
#define RALL(v) (v).rbegin(), (v).rend()
const double EPS = 1e-7;
const int INF = 1 << 30;
const ll LLINF = 1LL << 60;
constexpr int MOD = 1000000007;
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

//-------------------------------------

template <typename T> struct BIT {
  vector<T> v;
  BIT(int n) { v.assign(n + 1, 0); }
  // [0, k]
  T sum(int k) {
    T res = 0;
    for (++k; k > 0; k -= k & -k) {
      res += v[k];
    }
    return res;
  }
  // [l, r)
  T sum(int l, int r) {
    return (l == 0 ? sum(r - 1) : sum(r - 1) - sum(l - 1));
  }
  void add(int k, T x) {
    for (++k; k < v.size(); k += k & -k) {
      v[k] += x;
    }
  }
};

struct EulerTour {
  int n;
  vector<int> vs, dep, in, out;
  EulerTour(vector<vector<int>> &g, int _n)
      : n(_n), vs(2 * _n - 1), dep(2 * _n - 1), in(_n), out(_n) {
    build(g, 0);
  }
  void dfs(int v, int p, int d, int &k, vector<vector<int>> &g) {
    in[v] = k;
    vs[k] = v;
    dep[k++] = d;
    for (const int &u : g[v]) {
      if (u != p) {
        dfs(u, v, d + 1, k, g);
        vs[k] = v;
        dep[k++] = d;
      }
    }
  }
  void build(vector<vector<int>> &g, int root) {
    int k = 0;
    dfs(root, -1, 0, k, g);
    for (int i = 0; i < vs.size(); i++) {
      out[vs[i]] = i;
    }
  }
};

int N;
// colList[i] := 色iの頂点の集合
// idList[i] := 頂点iがeuler列に出現する時刻の集合
vector<vector<int>> G, colList, idList;
vector<int> c;

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cin >> N;
  G.resize(N);
  c.resize(N);
  for (int i = 0; i < N; i++) {
    cin >> c[i];
    c[i]--;
  }
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    G[a].emplace_back(b);
    G[b].emplace_back(a);
  }
  EulerTour euler(G, N);
  debugVector(euler.vs);
  int sz = euler.vs.size();
  BIT<ll> bit(sz);
  colList.resize(N);
  idList.resize(N);
  for (int i = 0; i < N; i++) {
    bit.add(euler.in[i], 1);
    colList[c[i]].emplace_back(i);
  }
  for (int i = 0; i < sz; i++) {
    idList[euler.vs[i]].emplace_back(i);
  }
  for (int i = 0; i < N; i++) {
    ll ans = (ll)N * (N + 1) / 2;
    vector<pair<int, int>> v;
    map<int, int> renketu;
    for (auto j : colList[i]) {
      bit.add(euler.in[j], -1);
      renketu[j] = 1;
      for (int id = 1; id < idList[j].size(); id++) {
        v.emplace_back(idList[j][id], j);
      }
    }
    sort(ALL(v));
    for (auto p : v) {
      ll now = bit.sum(euler.in[p.second], p.first);
      ans -= now * (now + 1) / 2;
      renketu[p.second] += now;
      bit.add(euler.in[p.second], -now);
    }
    {
      ll tmp = bit.sum(sz - 1);
      ans -= tmp * (tmp + 1) / 2;
    }
    cout << ans << "\\n";
    for (auto p : renketu) {
      bit.add(euler.in[p.first], p.second);
    }
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
using ll = long long;

#define ALL(v) (v).begin(), (v).end()
#define RALL(v) (v).rbegin(), (v).rend()
#define endl "\\n"
const double EPS = 1e-7;
const int INF = 1 << 30;
const ll LLINF = 1LL << 60;
const double PI = acos(-1);
constexpr int MOD = 1000000007;
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

//-------------------------------------

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n;
  cin >> n;
  vector<int> c(n);
  vector<vector<int>> g(n);
  for (int i = 0; i < n; i++) {
    cin >> c[i];
    c[i]--;
  }
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    g[--a].emplace_back(--b);
    g[b].emplace_back(a);
  }
  // sz[i] := 頂点iを根とする部分木のsize
  vector<int> sz(n, 0);
  // dp[i][c] := 頂点iを根とする部分木の中で、色cを
  // 根とする部分木のsizeの合計
  vector<map<int, int>> dp(n);
  vector<vector<int>> res(n);
  auto dfs = [&](auto &&dfs, int v, int p) -> void {
    for (const int &u : g[v]) {
      if (u != p) {
        dfs(dfs, u, v);
        sz[v] += sz[u];
        int tmp = sz[u];
        if (dp[u].count(c[v])) {
          tmp -= dp[u][c[v]];
        }
        res[c[v]].emplace_back(tmp);
        // if(dp[v].size() < dp[u].size()) {
        //     swap(dp[v], dp[u]);
        // }
        for (auto &p : dp[u]) {
          dp[v][p.first] += p.second;
        }
      }
    }
    sz[v]++;
    dp[v][c[v]] = sz[v];
  };
  dfs(dfs, 0, -1);
  for (int i = 0; i < n; i++) {
    res[i].emplace_back(n - dp[0][i]);
  }
  vector<ll> ans(n, (ll)n * (n + 1) / 2);
  for (int i = 0; i < n; i++) {
    for (auto &j : res[i]) {
      ans[i] -= (ll)j * (j + 1) / 2;
    }
  }
  for (const ll &i : ans) {
    cout << i << endl;
  }
}