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

#include <bits/stdc++.h>
using namespace std;
template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << "(" << p.first << "," << p.second << ")";
  return os;
}
#ifdef __LOCAL
#define debug(x) cerr << __LINE__ << ": " << #x << " = " << x << endl
#define debugArray(x, n)                                                       \\
  cerr << __LINE__ << ": " << #x << " = {";                                    \\
  for (long long hoge = 0; (hoge) < (n); ++(hoge))                             \\
    cerr << ((hoge) ? "," : "") << x[hoge];                                    \\
  cerr << "}" << endl
#else
#define debug(x) (void(0))
#define debugArray(x, n) (void(0))
#endif

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N;
  cin >> N;
  int c[N];
  for (int i = 0; i < N; i++)
    cin >> c[i], c[i]--;
  vector<vector<int>> graph(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    graph[a].push_back(b);
    graph[b].push_back(a);
  }
  int sz[N];
  fill(sz, sz + N, 1);
  long long ans[N];
  fill(ans, ans + N, (long long)N * (N + 1) / 2);
  map<int, int> mp[N];
  auto merge = [](map<int, int> &A, map<int, int> &B) {
    if (A.size() < B.size())
      swap(A, B);
    for (auto &p : B)
      A[p.first] += p.second;
    return A;
  };
  function<void(int, int)> dfs = [&](int v, int p) {
    for (int ch : graph[v])
      if (ch != p) {
        dfs(ch, v);
        sz[v] += sz[ch];
        if (mp[v].size() < mp[ch].size())
          swap(mp[v], mp[ch]);
        for (auto &p : mp[ch])
          mp[v][p.first] += p.second;
      }
    mp[v][c[v]] = sz[v];
    if (p != -1) {
      long long tmp = sz[v] - mp[v][c[p]];
      ans[c[p]] -= tmp * (tmp + 1) / 2;
    }
  };
  dfs(0, -1);
  for (int i = 0; i < N; i++) {
    long long tmp = N - mp[0][i];
    ans[i] -= tmp * (tmp + 1) / 2;
  }
  for (int i = 0; i < N; i++)
    cout << ans[i] << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << "(" << p.first << "," << p.second << ")";
  return os;
}
#ifdef __LOCAL
#define debug(x) cerr << __LINE__ << ": " << #x << " = " << x << endl
#define debugArray(x, n)                                                       \\
  cerr << __LINE__ << ": " << #x << " = {";                                    \\
  for (long long hoge = 0; (hoge) < (n); ++(hoge))                             \\
    cerr << ((hoge) ? "," : "") << x[hoge];                                    \\
  cerr << "}" << endl
#else
#define debug(x) (void(0))
#define debugArray(x, n) (void(0))
#endif

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N;
  cin >> N;
  int c[N];
  for (int i = 0; i < N; i++)
    cin >> c[i], c[i]--;
  vector<vector<int>> graph(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    graph[a].push_back(b);
    graph[b].push_back(a);
  }
  int sz[N];
  fill(sz, sz + N, 1);
  long long ans[N];
  fill(ans, ans + N, (long long)N * (N + 1) / 2);
  map<int, int> mp[N];
  auto merge = [](map<int, int> &A, map<int, int> &B) {
    if (A.size() < B.size())
      swap(A, B);
    for (auto &p : B)
      A[p.first] += p.second;
    return A;
  };
  function<void(int, int)> dfs = [&](int v, int p) {
    for (int ch : graph[v])
      if (ch != p) {
        dfs(ch, v);
        sz[v] += sz[ch];
        mp[v] = merge(mp[v], mp[ch]);
      }
    mp[v][c[v]] = sz[v];
    if (p != -1) {
      long long tmp = sz[v] - mp[v][c[p]];
      ans[c[p]] -= tmp * (tmp + 1) / 2;
    }
  };
  dfs(0, -1);
  for (int i = 0; i < N; i++) {
    long long tmp = N - mp[0][i];
    ans[i] -= tmp * (tmp + 1) / 2;
  }
  for (int i = 0; i < N; i++)
    cout << ans[i] << endl;
  return 0;
}