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

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#define popcount __builtin_popcount
using namespace std;
typedef long long int ll;
typedef pair<int, int> P;
vector<int> g[200020];
int n;
int c[200020], sz[200020];
ll s[200020];
void 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;
  }
}
void dfs(int x, int p, map<int, int> &mp) {
  sz[x] = 1;
  for (auto y : g[x]) {
    if (y == p)
      continue;
    map<int, int> mpy;
    dfs(y, x, mpy);
    sz[x] += sz[y];
    ll c1 = sz[y];
    if (mpy.find(c[x]) != mpy.end()) {
      c1 -= mpy[c[x]];
    }
    s[c[x]] += c1 * (c1 + 1) / 2;
    merge(mp, mpy);
  }
  mp[c[x]] = sz[x];
}
int main() {
  cin >> 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);
  }
  map<int, int> mp;
  dfs(0, -1, mp);
  for (int i = 0; i < n; i++) {
    if (i == c[0])
      continue;
    ll c1 = n - mp[i];
    s[i] += c1 * (c1 + 1) / 2;
  }
  for (int i = 0; i < n; i++)
    cout << (ll)n * (n + 1) / 2 - s[i] << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#define popcount __builtin_popcount
using namespace std;
typedef long long int ll;
typedef pair<int, int> P;
vector<int> g[200020];
int n;
int c[200020], sz[200020];
ll s[200020];
map<int, int> 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;
}
void dfs(int x, int p, map<int, int> &mp) {
  sz[x] = 1;
  for (auto y : g[x]) {
    if (y == p)
      continue;
    map<int, int> mpy;
    dfs(y, x, mpy);
    sz[x] += sz[y];
    ll c1 = sz[y];
    if (mpy.find(c[x]) != mpy.end()) {
      c1 -= mpy[c[x]];
    }
    s[c[x]] += c1 * (c1 + 1) / 2;
    mp = merge(mp, mpy);
  }
  mp[c[x]] = sz[x];
}
int main() {
  cin >> 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);
  }
  map<int, int> mp;
  dfs(0, -1, mp);
  for (int i = 0; i < n; i++) {
    if (i == c[0])
      continue;
    ll c1 = n - mp[i];
    s[i] += c1 * (c1 + 1) / 2;
  }
  for (int i = 0; i < n; i++)
    cout << (ll)n * (n + 1) / 2 - s[i] << endl;
  return 0;
}