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

#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#define all(x) (x).begin(), (x).end()
#define SZ(x) ((int)x.size())

using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using vi = vector<int>;
using vl = vector<ll>;

mt19937 mrand(random_device{}());
ll rnd(ll a, ll b) { return a + mrand() % (b - a + 1); }
ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }

const int N = 2e5 + 100;
vi g[N];
vl s[N];
int color[N], dp[N], cnt[N];
ll ans[N];
ll n;
map<int, int> per[N];

void merge(int i, int j) {
  if (per[i].size() < per[j].size())
    per[i].swap(per[j]);
  for (auto k : per[j]) {
    per[i][k.first] += k.second;
  }
  per[j].clear();
}

void dfs(int u, int p) {
  dp[u] = 1;
  for (int v : g[u]) {
    if (v == p)
      continue;
    dfs(v, u);
    dp[u] += dp[v];
    s[color[u]].push_back(dp[v] - per[v][color[u]]);
    merge(u, v);
  }
  per[u][color[u]] = dp[u];
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
#ifdef arrias
  freopen("a.txt", "r", stdin);
#endif
  cin >> n;
  rep(i, 1, n) cin >> color[i], cnt[color[i]]++;
  rep(i, 1, n - 1) {
    int u, v;
    cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs(1, -1);
  rep(i, 1, n) {
    if (cnt[i] == 0) {
      cout << 0 << "\\n";
      continue;
    }
    ans[i] = n * (n - 1) / 2 + n;
    ll sum = accumulate(all(s[i]), 0ll);
    if (sum + cnt[i] != n) {
      s[i].push_back(n - cnt[i] - sum);
    }
    for (auto j : s[i]) {
      ans[i] -= ((j * (j - 1) / 2) + j);
    }
    cout << ans[i] << "\\n";
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#define all(x) (x).begin(), (x).end()
#define SZ(x) ((int)x.size())

using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using vi = vector<int>;
using vl = vector<ll>;

mt19937 mrand(random_device{}());
ll rnd(ll a, ll b) { return a + mrand() % (b - a + 1); }
ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }

const int N = 2e5 + 100;
vi g[N];
vl s[N];
int color[N], dp[N], cnt[N];
ll ans[N];
ll n;

void fromSmallToLarge(map<int, int> &a, map<int, int> &b) {
  if (SZ(a) > SZ(b))
    swap(a, b);
  for (auto i : b)
    a[i.first] += i.second;
}

map<int, int> dfs(int u, int p) {
  dp[u] = 1;
  map<int, int> ret;
  for (int v : g[u]) {
    if (v == p)
      continue;
    auto x = dfs(v, u);
    dp[u] += dp[v];
    s[color[u]].push_back(dp[v] - x[color[u]]);
    fromSmallToLarge(ret, x);
  }
  ret[color[u]] = dp[u];
  return ret;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
#ifdef arrias
  freopen("a.txt", "r", stdin);
#endif
  cin >> n;
  rep(i, 1, n) cin >> color[i], cnt[color[i]]++;
  rep(i, 1, n - 1) {
    int u, v;
    cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs(1, -1);
  rep(i, 1, n) {
    if (cnt[i] == 0) {
      cout << 0 << "\\n";
      continue;
    }
    ans[i] = n * (n - 1) / 2 + n;
    ll sum = accumulate(all(s[i]), 0ll);
    if (sum + cnt[i] != n) {
      s[i].push_back(n - cnt[i] - sum);
    }
    for (auto j : s[i]) {
      ans[i] -= ((j * (j - 1) / 2) + j);
    }
    cout << ans[i] << "\\n";
  }
  return 0;
}