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

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); i++)
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;

int n;

vector<int> E[200000];
bool used[200000];
int sz[200000];
ll ans[200000];
int c[300000];

int N, s;

#define map unordered_map

void szdfs(int v, int p) {
  sz[v] = 1;
  for (auto u : E[v]) {
    if (u == p || used[u])
      continue;
    szdfs(u, v);
    sz[v] += sz[u];
  }
}

void search_centroid(int v, int p) {
  bool centroid = true;
  for (auto u : E[v]) {
    if (u == p || used[u])
      continue;
    search_centroid(u, v);
    if (sz[u] > N / 2)
      centroid = false;
  }
  if (N - sz[v] > N / 2)
    centroid = false;
  if (centroid)
    s = v;
}

int S[200000];
int M[200000];
void pathdfs(int v, int p, map<int, ll> &mp, ll &X) {
  S[v] = 1;
  X++;
  bool b = (M[c[v]] == 0);
  M[c[v]]++;
  for (auto u : E[v]) {
    if (u == p || used[u])
      continue;
    pathdfs(u, v, mp, X);
    S[v] += S[u];
  }
  if (b)
    mp[c[v]] += S[v];
  M[c[v]]--;
}

void count_pairs(map<int, ll> &mp, ll X, ll x) {
  for (auto p : mp) {
    ans[p.first] +=
        p.second * (X - p.second) * x + p.second * (p.second - 1) / 2 * x;
    if (x > 0)
      ans[p.first] += p.second;
  }
}

void solve(int v) {
  szdfs(v, -1);
  N = sz[v];
  search_centroid(v, -1);
  used[s] = true;
  map<int, ll> mp_all;
  ll X_all = 0;
  for (auto u : E[s]) {
    if (used[u])
      continue;
    map<int, ll> mp;
    ll X = 0;
    mp[c[s]] = 0;
    M[c[s]]++;
    pathdfs(u, s, mp, X);
    M[c[s]]--;
    mp[c[s]] += S[u];
    count_pairs(mp, X, -1);
    X_all += X;
    if (mp_all.size() < mp.size())
      swap(mp, mp_all);
    for (auto p : mp)
      mp_all[p.first] += p.second;
  }
  ans[c[s]]++;
  count_pairs(mp_all, X_all, 1);
  for (auto u : E[s]) {
    if (!used[u])
      solve(u);
  }
}

int main() {
  int n;
  cin >> n;
  rep(i, n) scanf("%d", &c[i]), c[i]--;
  rep(i, n - 1) {
    int a, b;
    scanf("%d%d", &a, &b);
    a--;
    b--;
    E[a].push_back(b);
    E[b].push_back(a);
  }
  solve(0);
  rep(i, n) { printf("%lld\\n", ans[i]); }
}
/*
6
5 3 3 1 6 6
1 4
2 4
2 5
2 6
3 4
*/

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); i++)
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;

int n;

vector<int> E[200000];
bool used[200000];
int sz[200000];
ll ans[200000];
int c[300000];

int N, s;

#define map unordered_map

void szdfs(int v, int p) {
  sz[v] = 1;
  for (auto u : E[v]) {
    if (u == p || used[u])
      continue;
    szdfs(u, v);
    sz[v] += sz[u];
  }
}

void search_centroid(int v, int p) {
  bool centroid = true;
  for (auto u : E[v]) {
    if (u == p || used[u])
      continue;
    search_centroid(u, v);
    if (sz[u] > N / 2)
      centroid = false;
  }
  if (N - sz[v] > N / 2)
    centroid = false;
  if (centroid)
    s = v;
}

int S[200000];
map<int, int> M;
void pathdfs(int v, int p, map<int, ll> &mp, ll &X) {
  S[v] = 1;
  X++;
  bool b = (M[c[v]] == 0);
  M[c[v]]++;
  for (auto u : E[v]) {
    if (u == p || used[u])
      continue;
    pathdfs(u, v, mp, X);
    S[v] += S[u];
  }
  if (b)
    mp[c[v]] += S[v];
  M[c[v]]--;
}

void count_pairs(map<int, ll> &mp, ll X, ll x) {
  for (auto p : mp) {
    ans[p.first] +=
        p.second * (X - p.second) * x + p.second * (p.second - 1) / 2 * x;
    if (x > 0)
      ans[p.first] += p.second;
  }
}

void solve(int v) {
  szdfs(v, -1);
  N = sz[v];
  search_centroid(v, -1);
  used[s] = true;
  map<int, ll> mp_all;
  ll X_all = 0;
  for (auto u : E[s]) {
    if (used[u])
      continue;
    map<int, ll> mp;
    ll X = 0;
    mp[c[s]] = 0;
    M.clear();
    M[c[s]]++;
    pathdfs(u, s, mp, X);
    mp[c[s]] += S[u];
    count_pairs(mp, X, -1);
    X_all += X;
    if (mp_all.size() < mp.size())
      swap(mp, mp_all);
    for (auto p : mp)
      mp_all[p.first] += p.second;
  }
  ans[c[s]]++;
  count_pairs(mp_all, X_all, 1);
  for (auto u : E[s]) {
    if (!used[u])
      solve(u);
  }
}

int main() {
  int n;
  cin >> n;
  rep(i, n) scanf("%d", &c[i]), c[i]--;
  rep(i, n - 1) {
    int a, b;
    scanf("%d%d", &a, &b);
    a--;
    b--;
    E[a].push_back(b);
    E[b].push_back(a);
  }
  solve(0);
  rep(i, n) { printf("%lld\\n", ans[i]); }
}
/*
6
5 3 3 1 6 6
1 4
2 4
2 5
2 6
3 4
*/