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

#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;

V c;
Vll ans, sz, dp;

using Graph = vector<vector<int>>;
Graph G;

void dfs(int u, int p = -1) {
  ll tmp = dp[c[u]];
  for (auto v : G[u]) {
    if (v == p)
      continue;
    dp[c[u]] = 0;
    dfs(v, u);
    sz[u] += sz[v];
    ll x = sz[v] - dp[c[u]];
    ans[c[u]] -= x * (x + 1) / 2;
  }
  dp[c[u]] = tmp + sz[u];
}

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

  c.assign(n, 0);
  sz.assign(n, 1);
  dp.assign(n, 0);
  ans.assign(n, n * (n + 1) / 2);

  rep(i, n) {
    cin >> c[i];
    --c[i];
  }

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

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

  cout << ans;
  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;

V c;
Vll ans;

using Graph = vector<vector<int>>;
Graph G;

template <typename T> T my_merge(T a, T b) {
  if (a.size() < b.size())
    a.swap(b);
  for (auto &e : b)
    a[e.first] += e.second;
  return a;
}

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;
    ll sz_v;
    map<int, ll> dp;
    tie(sz_v, dp) = dfs(v, u);
    ll x = sz_v - dp[c[u]];
    ans[c[u]] -= x * (x + 1) / 2;
    res = my_merge(res, dp);
    sz += sz_v;
  }
  res[c[u]] = sz;
  return {sz, res};
}

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

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

  rep(i, n) {
    int x;
    cin >> x;
    c[i] = x - 1;
  }

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

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

  cout << ans;
  return 0;
}