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

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(x) (x).begin(), (x).end()
#define ln '\\n'
constexpr long long MOD = 1000000007LL;
// constexpr long long MOD = 998244353LL;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pii;
typedef pair<long long, long long> pll;
template <class T, class U> inline bool chmax(T &a, U b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T, class U> inline bool chmin(T &a, U b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

vector<int> G[201010];
int c[201010];
ll ans[201010];
ll sub[201010];
unordered_map<int, vector<int>> dic[201010];
void dfs(int v, int pv) {
  for (auto nv : G[v]) {
    if (nv == pv)
      continue;
    dfs(nv, v);
    sub[v] += sub[nv];
    if (dic[nv].find(c[v]) == dic[nv].end())
      ans[c[v]] += (sub[nv]) * (sub[nv] + 1) / 2;
    else {
      ll val = 0;
      for (auto j : dic[nv][c[v]]) {
        val += sub[j];
      }
      ans[c[v]] += (sub[nv] - val) * (sub[nv] - val + 1) / 2;
    }
    if (dic[v].size() < dic[nv].size()) {
      swap(dic[v], dic[nv]);
    }
    for (auto A : dic[nv]) {
      for (auto j : A.second) {
        dic[v][A.first].emplace_back(j);
      }
    }
  }
  dic[v][c[v]].clear();
  dic[v][c[v]].emplace_back(v);
  sub[v]++;
  return;
}
int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  ll N;
  cin >> N;
  set<int> zzz;
  rep(i, N) {
    cin >> c[i];
    --c[i];
    zzz.emplace(c[i]);
  }
  rep(i, N - 1) {
    int u, v;
    cin >> u >> v;
    --u;
    --v;
    G[u].emplace_back(v);
    G[v].emplace_back(u);
  }
  dfs(0, -1);
  // rep(i,N) cout << ans[i] << ln;
  for (auto i : dic[0]) {
    ll val = 0;
    for (auto j : i.second) {
      val += sub[j];
    }
    ans[i.first] += (sub[0] - val) * (sub[0] - val + 1) / 2;
  }
  rep(i, N) {
    if (!zzz.count(i))
      cout << 0 << ln;
    else
      cout << (N + 1) * N / 2 - ans[i] << ln;
  }
}


Error: TLE

Fixed submission:

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(x) (x).begin(), (x).end()
#define ln '\\n'
constexpr long long MOD = 1000000007LL;
// constexpr long long MOD = 998244353LL;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pii;
typedef pair<long long, long long> pll;
template <class T, class U> inline bool chmax(T &a, U b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T, class U> inline bool chmin(T &a, U b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

vector<int> G[201010];
int c[201010];
ll ans[201010];
ll sub[201010];
gp_hash_table<int, vector<int>> dic[201010];
void dfs(int v, int pv) {
  for (auto nv : G[v]) {
    if (nv == pv)
      continue;
    dfs(nv, v);
    sub[v] += sub[nv];
    if (dic[nv].find(c[v]) == dic[nv].end())
      ans[c[v]] += (sub[nv]) * (sub[nv] + 1) / 2;
    else {
      ll val = 0;
      for (auto j : dic[nv][c[v]]) {
        val += sub[j];
      }
      ans[c[v]] += (sub[nv] - val) * (sub[nv] - val + 1) / 2;
    }
    if (dic[v].size() < dic[nv].size()) {
      swap(dic[v], dic[nv]);
    }
    for (auto A : dic[nv]) {
      for (auto j : A.second) {
        dic[v][A.first].emplace_back(j);
      }
    }
  }
  dic[v][c[v]].clear();
  dic[v][c[v]].emplace_back(v);
  sub[v]++;
  return;
}
int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  ll N;
  cin >> N;
  set<int> zzz;
  rep(i, N) {
    cin >> c[i];
    --c[i];
    zzz.emplace(c[i]);
  }
  rep(i, N - 1) {
    int u, v;
    cin >> u >> v;
    --u;
    --v;
    G[u].emplace_back(v);
    G[v].emplace_back(u);
  }
  dfs(0, -1);
  // rep(i,N) cout << ans[i] << ln;
  for (auto i : dic[0]) {
    ll val = 0;
    for (auto j : i.second) {
      val += sub[j];
    }
    ans[i.first] += (sub[0] - val) * (sub[0] - val + 1) / 2;
  }
  rep(i, N) {
    if (!zzz.count(i))
      cout << 0 << ln;
    else
      cout << (N + 1) * N / 2 - ans[i] << ln;
  }
}