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

#include <bits/stdc++.h>
#define FOR(i, k, n) for (int i = (k); i < (n); ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) begin(x), end(x)

using namespace std;
using ll = int64_t;
using vecint = vector<int>;
using vecll = vector<ll>;

pair<map<int, ll>, int> dfs(const vector<vecint> &g, const vecint &clr,
                            vecll &ans, int i, int p = -1) {
  ll sz = 1;
  map<int, ll> ret;
  for (auto &&to : g[i]) {
    if (to == p)
      continue;
    auto [cret, csz] = dfs(g, clr, ans, to, i);
    ll icsz = csz - cret[clr[i]];
    if (cret.size() > ret.size())
      swap(cret, ret);
    for (auto &&[c, num] : cret) {
      ret[c] += num;
    }
    ans[clr[i]] -= icsz * (icsz + 1) / 2;
    sz += csz;
  }
  ret[clr[i]] = sz;
  return make_pair(move(ret), sz);
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  ll n;
  cin >> n;
  vecint clr(n);
  REP(i, n) cin >> clr[i];
  vector<vecint> g(n);
  REP(i, n - 1) {
    int a, b;
    cin >> a >> b;
    --a;
    --b;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  vecll ans(n + 1, n * (n + 1) / 2);
  auto [cret, _] = dfs(g, clr, ans, 0);
  FOR(i, 1, n + 1) {
    ll icsz = n - cret[i];
    ans[i] -= icsz * (icsz + 1) / 2;
    cout << ans[i] << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define _GLIBCXX_DEBUG
#define rep2(i, m, n) for (int i = int(m); i < int(n); ++i)
#define drep2(i, m, n) for (int i = int(m - 1); i >= int(n); --i)
#define rep(i, n) rep2(i, 0, n)
#define drep(i, n) drep2(i, n, 0)
#define all(a) a.begin(), a.end()
using ll = long long;
using ld = long double;
using V = vector<int>;
using Vll = vector<ll>;
using Vld = vector<ld>;
using Vbo = vector<bool>;
using VV = vector<V>;
using VVll = vector<Vll>;
using VVld = vector<Vld>;
using VVbo = vector<Vbo>;
using P = pair<int, int>;
using Pll = pair<ll, ll>;
using Pld = pair<ld, ld>;
struct fast_ios {
  fast_ios() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);
  };
} fast_ios_;
template <typename T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  os << "(" << p.first << ", " << p.second << ")";
  return os;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &e : v)
    is >> e;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (auto &e : v)
    os << e << endl;
  return os;
}
template <typename T> inline int count_between(vector<T> &a, T l, T r) {
  return lower_bound(all(a), r) - lower_bound(all(a), l);
} // [l, r)
inline int Log2(ll x) {
  int k;
  for (k = 0; x > 0; ++k)
    x >>= 1;
  return k;
} // number of binary digits
const int INF = 1 << 30;
const ll INFll = 1ll << 62;
const ld EPS = 1e-10;
const int MOD = int(1e9) + 7;

VV G;
V c;
Vll ans;

pair<ll, map<int, ll>> dfs(int u, int p = -1) {
  map<int, ll> res;
  ll sz = 1;
  for (auto v : G[u]) {
    if (v == p)
      continue;
    auto [sz_v, res_v] = dfs(v, u);
    ll x = sz_v - res_v[c[u]];
    ans[c[u]] -= x * (x + 1) / 2;
    if (res.size() < res_v.size())
      swap(res, res_v);
    for (auto &[a, b] : res_v)
      res[a] += b;
    sz += sz_v;
  }
  res[c[u]] = sz;
  return {sz, res};
}

int main() {
  ll n;
  cin >> n;

  G.assign(n, V());
  c.assign(n, 0);
  ans.assign(n, n * (n + 1) / 2);

  rep(i, n) {
    cin >> c[i];
    --c[i];
  }
  rep(i, n - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  auto [sz, res] = dfs(0);
  rep(i, n) {
    ll x = sz - res[i];
    ans[i] -= x * (x + 1) / 2;
  }

  cout << ans;
  return 0;
}