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

#include <algorithm>
#include <iostream>
#include <map>
#include <vector>

using namespace std;

long long search(const vector<vector<int>> &g, const vector<int> &col, int pos,
                 int prev, vector<long long> &dec,
                 vector<map<int, long long>> &mp, vector<int> &ref) {
  int cnt = 1;
  for (auto t : g[pos]) {
    if (t == prev)
      continue;
    auto res = search(g, col, t, pos, dec, mp, ref);
    auto &oppMp = mp[ref[t]];
    auto num = res;
    if (oppMp.count(col[pos]))
      num -= oppMp[col[pos]];
    dec[col[pos]] += num * (num + 1) / 2;
    if (mp[ref[pos]].size() < mp[ref[t]].size())
      swap(ref[pos], ref[t]);
    for (auto &it : mp[ref[t]]) {
      if (!mp[ref[pos]].count(it.first))
        mp[ref[pos]][it.first] = 0;
      mp[ref[pos]][it.first] += it.second;
    }
    cnt += res;
  }
  mp[ref[pos]][col[pos]] = cnt;
  if (prev == -1) {
    for (int i = 0; i < dec.size(); i++) {
      long long num = dec.size();
      if (mp[ref[pos]].count(i))
        num -= mp[ref[pos]][i];
      dec[i] += num * (num + 1) / 2;
    }
  }
  return cnt;
}

int main() {
  long long N;
  cin >> N;
  vector<int> c(N);
  for (auto &t : c) {
    cin >> t;
    --t;
  }
  vector<vector<int>> g(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    cin >> a >> b;
    --a;
    --b;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  vector<long long> dec(N, 0);
  vector<map<int, long long>> mp(N);
  vector<int> ref(N);
  for (int i = 0; i < N; i++)
    ref[i] = i;
  search(g, c, 0, -1, dec, mp, ref);
  for (int i = 0; i < N; i++) {
    cout << N * (N + 1) / 2 - dec[i] << endl;
  }
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <map>
#include <vector>

using namespace std;

pair<int, map<int, long long>> search(const vector<vector<int>> &g,
                                      const vector<int> &col, int pos, int prev,
                                      vector<long long> &dec) {
  int cnt = 1;
  map<int, long long> mp;
  for (auto t : g[pos]) {
    if (t == prev)
      continue;
    auto res = search(g, col, t, pos, dec);
    auto num = res.second.count(col[pos]) ? res.first - res.second[col[pos]]
                                          : res.first;
    dec[col[pos]] += num * (num + 1) / 2;
    if (mp.size() < res.second.size())
      swap(mp, res.second);
    for (auto &it : res.second) {
      if (!mp.count(it.first))
        mp[it.first] = 0;
      mp[it.first] += it.second;
    }
    cnt += res.first;
  }
  mp[col[pos]] = cnt;
  if (prev == -1) {
    for (int i = 0; i < dec.size(); i++) {
      long long num = dec.size();
      if (mp.count(i))
        num -= mp[i];
      dec[i] += num * (num + 1) / 2;
    }
  }
  return make_pair(cnt, mp);
}

int main() {
  long long N;
  cin >> N;
  vector<int> c(N);
  for (auto &t : c) {
    cin >> t;
    --t;
  }
  vector<vector<int>> g(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    cin >> a >> b;
    --a;
    --b;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  vector<long long> dec(N, 0);
  search(g, c, 0, -1, dec);
  for (int i = 0; i < N; i++) {
    cout << N * (N + 1) / 2 - dec[i] << endl;
  }
}