Problem p02710 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

using namespace std;
#define int long long
const int MOD = 1000000007;
vector<int> C;
vector<vector<int>> g;
vector<int> res;
vector<int> sz;
void dfs(int a, int p) {
  sz[a] = 1;
  for (int i : g[a]) {
    if (i != p) {
      dfs(i, a);
      sz[a] += sz[i];
    }
  }
}
int dfs2(int a, int p, map<int, int> &mp, int &add) {
  int mx = -1;
  int mxn = -1;
  for (int i : g[a]) {
    if (i != p && mx < sz[i]) {
      mx = sz[i];
      mxn = i;
    }
  }
  if (mxn != -1) {
    dfs2(mxn, a, mp, add);
    {
      int t = mp[C[a]] + add;
      res[C[a]] += t * (t + 1) / 2;
    }
    for (int i : g[a]) {
      if (i != p && i != mxn) {
        map<int, int> mp2;
        int add2 = 0;
        dfs2(i, a, mp2, add2);
        {
          int t = mp2[C[a]] + add2;
          res[C[a]] += t * (t + 1) / 2;
        }
        add += add2;
        for (auto m : mp2) {
          mp[m.first] += m.second;
        }
      }
    }
  }
  add++;
  mp[C[a]] = -add;
}
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N;
  cin >> N;
  C.resize(N);
  g.resize(N);
  res.resize(N, 0);
  sz.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].push_back(b);
    g[b].push_back(a);
  }
  dfs(0, -1);
  map<int, int> mp;
  int add = 0;
  dfs2(0, -1, mp, add);
  for (int i = 0; i < N; i++) {
    int t = mp[i] + add;
    res[i] += t * (t + 1) / 2;
  }
  int all = N * (N + 1) / 2;
  for (int i = 0; i < N; i++) {
    cout << all - res[i] << endl;
  }
}

Fixed submission:
#include <bits/stdc++.h>

using namespace std;
#define int long long
const int MOD = 1000000007;
vector<int> C;
vector<vector<int>> g;
vector<int> res;
vector<int> sz;
void dfs(int a, int p) {
  sz[a] = 1;
  for (int i : g[a]) {
    if (i != p) {
      dfs(i, a);
      sz[a] += sz[i];
    }
  }
}
void dfs2(int a, int p, map<int, int> &mp, int &add) {
  int mx = -1;
  int mxn = -1;
  for (int i : g[a]) {
    if (i != p && mx < sz[i]) {
      mx = sz[i];
      mxn = i;
    }
  }
  if (mxn != -1) {
    dfs2(mxn, a, mp, add);
    {
      int t = mp[C[a]] + add;
      res[C[a]] += t * (t + 1) / 2;
    }
    for (int i : g[a]) {
      if (i != p && i != mxn) {
        map<int, int> mp2;
        int add2 = 0;
        dfs2(i, a, mp2, add2);
        {
          int t = mp2[C[a]] + add2;
          res[C[a]] += t * (t + 1) / 2;
        }
        add += add2;
        for (auto m : mp2) {
          mp[m.first] += m.second;
        }
      }
    }
  }
  add++;
  mp[C[a]] = -add;
}
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N;
  cin >> N;
  C.resize(N);
  g.resize(N);
  res.resize(N, 0);
  sz.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].push_back(b);
    g[b].push_back(a);
  }
  dfs(0, -1);
  map<int, int> mp;
  int add = 0;
  dfs2(0, -1, mp, add);
  for (int i = 0; i < N; i++) {
    int t = mp[i] + add;
    res[i] += t * (t + 1) / 2;
  }
  int all = N * (N + 1) / 2;
  for (int i = 0; i < N; i++) {
    cout << all - res[i] << endl;
  }
}