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

// #define _GLIBCXX_DEBUG

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

#define endl '\\n'
#define lfs cout << fixed << setprecision(10)
#define ALL(a) (a).begin(), (a).end()
#define ALLR(a) (a).rbegin(), (a).rend()
#define spa << " " <<
#define fi first
#define se second
#define MP make_pair
#define MT make_tuple
#define PB push_back
#define EB emplace_back
#define rep(i, n, m) for (ll i = (n); i < (ll)(m); i++)
#define rrep(i, n, m) for (ll i = (ll)(m)-1; i >= (ll)(n); i--)
using ll = long long;
using ld = long double;
const ll MOD1 = 1e9 + 7;
const ll MOD9 = 998244353;
const ll INF = 1e18;
using P = pair<ll, ll>;
template <typename T1, typename T2> bool chmin(T1 &a, T2 b) {
  if (a > b) {
    a = b;
    return true;
  } else
    return false;
}
template <typename T1, typename T2> bool chmax(T1 &a, T2 b) {
  if (a < b) {
    a = b;
    return true;
  } else
    return false;
}
ll median(ll a, ll b, ll c) {
  return a + b + c - max({a, b, c}) - min({a, b, c});
}
void ans1(bool x) {
  if (x)
    cout << "Yes" << endl;
  else
    cout << "No" << endl;
}
void ans2(bool x) {
  if (x)
    cout << "YES" << endl;
  else
    cout << "NO" << endl;
}
void ans3(bool x) {
  if (x)
    cout << "Yay!" << endl;
  else
    cout << ":(" << endl;
}
template <typename T1, typename T2> void ans(bool x, T1 y, T2 z) {
  if (x)
    cout << y << endl;
  else
    cout << z << endl;
}
template <typename T> void debug(vector<vector<T>> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    cout << v[i][0];
    for (ll j = 1; j < w; j++)
      cout spa v[i][j];
    cout << endl;
  }
};
void debug(vector<string> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    for (ll j = 0; j < w; j++)
      cout << v[i][j];
    cout << endl;
  }
};
template <typename T> void debug(vector<T> &v, ll n) {
  if (n != 0)
    cout << v[0];
  for (ll i = 1; i < n; i++)
    cout spa v[i];
  cout << endl;
};
template <typename T> vector<vector<T>> vec(ll x, ll y, T w) {
  vector<vector<T>> v(x, vector<T>(y, w));
  return v;
}
ll gcd(ll x, ll y) {
  ll r;
  while (y != 0 && (r = x % y) != 0) {
    x = y;
    y = r;
  }
  return y == 0 ? x : y;
}
vector<ll> dx = {1, 0, -1, 0, 1, 1, -1, -1};
vector<ll> dy = {0, 1, 0, -1, 1, -1, 1, -1};
template <typename T> vector<T> make_v(size_t a, T b) {
  return vector<T>(a, b);
}
template <typename... Ts> auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v(ts...))>(a, make_v(ts...));
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  return os << p.first << " " << p.second;
}
vector<ll> col(200005);
vector<ll> sz(200005, 1);
void merge(unordered_map<ll, ll> &a, unordered_map<ll, ll> &b) {
  if (a.size() < b.size())
    swap(a, b);
  for (auto z : b)
    a[z.fi] += z.se;
}
vector<vector<ll>> comp(200005);
unordered_map<ll, ll> dfs(ll k, vector<vector<ll>> &g, ll par) {
  unordered_map<ll, ll> ret;
  for (auto to : g[k]) {
    if (to != par) {
      auto mp = dfs(to, g, k);
      sz[k] += sz[to];
      // cout<<col[k] spa sz[to] spa mp[col[k]]<<endl;
      comp[col[k]].PB(sz[to] - mp[col[k]]);
      mp[col[k]] = 0;
      merge(ret, mp);
    }
  }
  // cout<<sz[k]<<endl;
  ret[col[k]] = sz[k];
  return ret;
}
int main() {
  cin.tie(nullptr);
  ios_base::sync_with_stdio(false);
  ll res = 0, buf = 0;
  bool judge = true;
  ll n;
  cin >> n;
  rep(i, 0, n) {
    cin >> col[i];
    col[i]--;
  }
  vector<vector<ll>> g(n);
  rep(i, 0, n - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    g[a].PB(b);
    g[b].PB(a);
  }
  auto mp = dfs(0, g, -1);
  rep(i, 0, n) {
    // cout<<mp[i]<<endl;
    comp[i].PB(n - mp[i]);
  }
  rep(i, 0, n) {
    // debug(comp[i],comp[i].size());
    ll ret = n * (n + 1) / 2;
    for (auto z : comp[i])
      ret -= z * (z + 1) / 2;
    cout << ret << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

// #define _GLIBCXX_DEBUG

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

#define endl '\\n'
#define lfs cout << fixed << setprecision(10)
#define ALL(a) (a).begin(), (a).end()
#define ALLR(a) (a).rbegin(), (a).rend()
#define spa << " " <<
#define fi first
#define se second
#define MP make_pair
#define MT make_tuple
#define PB push_back
#define EB emplace_back
#define rep(i, n, m) for (ll i = (n); i < (ll)(m); i++)
#define rrep(i, n, m) for (ll i = (ll)(m)-1; i >= (ll)(n); i--)
using ll = long long;
using ld = long double;
const ll MOD1 = 1e9 + 7;
const ll MOD9 = 998244353;
const ll INF = 1e18;
using P = pair<ll, ll>;
template <typename T1, typename T2> bool chmin(T1 &a, T2 b) {
  if (a > b) {
    a = b;
    return true;
  } else
    return false;
}
template <typename T1, typename T2> bool chmax(T1 &a, T2 b) {
  if (a < b) {
    a = b;
    return true;
  } else
    return false;
}
ll median(ll a, ll b, ll c) {
  return a + b + c - max({a, b, c}) - min({a, b, c});
}
void ans1(bool x) {
  if (x)
    cout << "Yes" << endl;
  else
    cout << "No" << endl;
}
void ans2(bool x) {
  if (x)
    cout << "YES" << endl;
  else
    cout << "NO" << endl;
}
void ans3(bool x) {
  if (x)
    cout << "Yay!" << endl;
  else
    cout << ":(" << endl;
}
template <typename T1, typename T2> void ans(bool x, T1 y, T2 z) {
  if (x)
    cout << y << endl;
  else
    cout << z << endl;
}
template <typename T> void debug(vector<vector<T>> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    cout << v[i][0];
    for (ll j = 1; j < w; j++)
      cout spa v[i][j];
    cout << endl;
  }
};
void debug(vector<string> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    for (ll j = 0; j < w; j++)
      cout << v[i][j];
    cout << endl;
  }
};
template <typename T> void debug(vector<T> &v, ll n) {
  if (n != 0)
    cout << v[0];
  for (ll i = 1; i < n; i++)
    cout spa v[i];
  cout << endl;
};
template <typename T> vector<vector<T>> vec(ll x, ll y, T w) {
  vector<vector<T>> v(x, vector<T>(y, w));
  return v;
}
ll gcd(ll x, ll y) {
  ll r;
  while (y != 0 && (r = x % y) != 0) {
    x = y;
    y = r;
  }
  return y == 0 ? x : y;
}
vector<ll> dx = {1, 0, -1, 0, 1, 1, -1, -1};
vector<ll> dy = {0, 1, 0, -1, 1, -1, 1, -1};
template <typename T> vector<T> make_v(size_t a, T b) {
  return vector<T>(a, b);
}
template <typename... Ts> auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v(ts...))>(a, make_v(ts...));
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  return os << p.first << " " << p.second;
}
vector<ll> col(200005);
vector<ll> sz(200005, 1);
void merge(map<ll, ll> &a, map<ll, ll> &b) {
  // if(a.size()<b.size())swap(a,b);
  for (auto z : b)
    a[z.fi] += z.se;
}
vector<vector<ll>> comp(200005);
map<ll, ll> dfs(ll k, vector<vector<ll>> &g, ll par) {
  map<ll, ll> ret;
  for (auto to : g[k]) {
    if (to != par) {
      auto mp = dfs(to, g, k);
      sz[k] += sz[to];
      // cout<<col[k] spa sz[to] spa mp[col[k]]<<endl;
      comp[col[k]].PB(sz[to] - mp[col[k]]);
      mp[col[k]] = 0;
      merge(ret, mp);
    }
  }
  // cout<<sz[k]<<endl;
  ret[col[k]] = sz[k];
  return ret;
}
int main() {
  cin.tie(nullptr);
  ios_base::sync_with_stdio(false);
  ll res = 0, buf = 0;
  bool judge = true;
  ll n;
  cin >> n;
  rep(i, 0, n) {
    cin >> col[i];
    col[i]--;
  }
  vector<vector<ll>> g(n);
  rep(i, 0, n - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    g[a].PB(b);
    g[b].PB(a);
  }
  auto mp = dfs(0, g, -1);
  rep(i, 0, n) {
    // cout<<mp[i]<<endl;
    comp[i].PB(n - mp[i]);
  }
  rep(i, 0, n) {
    // debug(comp[i],comp[i].size());
    ll ret = n * (n + 1) / 2;
    for (auto z : comp[i])
      ret -= z * (z + 1) / 2;
    cout << ret << endl;
  }
  return 0;
}