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

#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<double, double> pdd;
const ll mod = 1e9 + 7;
// const ll mod = 998244353;
#define REP(i, n) for (int i = 0; i < (int)n; ++i)
// debug
#define dump(x) cerr << #x << " = " << (x) << endl;
#define debug(x)                                                               \\
  cerr << #x << " = " << (x) << " (L" << __LINE__ << ")"                       \\
       << " " << __FILE__ << endl;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> v) {
  os << "(" << v.first << ", " << v.second << ")";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << " ";
    }
    os << v[i];
  }
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<vector<T>> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << endl;
    }
    os << v[i];
  }
  return os;
}
string num2bit(ll num, ll len) {
  string bit = "";
  REP(i, len) { bit += char('0' + (num >> i & 1)); }
  return bit;
}

vector<ll> c(202020, 0);
using Graph = vector<vector<int>>;
using pml = pair<map<ll, ll>, ll>;

void printmp(map<ll, ll> &mp) {
  for (auto x : mp) {
    cerr << x << " ";
  }
  cerr << endl;
}

ll N;
ll kazu = 0;
vector<ll> cnt(202020, 0);
vector<ll> sz(202020, 0);
map<ll, ll> dfs(ll now, ll par, Graph &G) {
  sz[now] = 1;
  map<ll, ll> mp;
  for (auto next : G[now]) {
    if (next == par)
      continue;
    auto mpch = dfs(next, now, G);

    sz[now] += sz[next];
    // ll val = szch - mpch[c[now]];
    ll val = sz[next];
    if (mpch.find(c[now]) != mpch.end())
      val -= mpch[c[now]];
    cnt[c[now]] -= val * (val + 1) / 2;

    if ((int)mp.size() < (int)mpch.size())
      swap(mp, mpch);

    /*
    cout << "mp: ";
    printmp(mp);
    cout << "mpch: ";
    printmp(mpch);
    cout << endl;
    */

    for (auto x : mpch) {
      mp[x.first] += x.second;
      kazu++;
    }
    mpch.clear();
  }

  /*
  dump(now)
  printmp(mp);
  dump(sz)
  dump(mp[c[now]])
  */

  // ll val = sz - 1 - mp[c[now]];
  // cnt[c[now]] -= val*(val+1)/2;
  mp[c[now]] = sz[now];

  /*
  REP(i, N){
      cout << cnt[i] << " ";
  }
  cout << endl;
  cout << endl;
  */
  return mp;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  cin >> N;

  REP(i, N) cin >> c[i], c[i]--;

  Graph G(N);
  REP(i, N - 1) {
    ll a, b;
    cin >> a >> b;
    a--, b--;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  auto mp = dfs(0, -1, G);

  // printmp(mp);
  vector<ll> nokori(N, N);

  for (auto x : mp) {
    nokori[x.first] = N - x.second;
  }
  // cout << nokori << endl;
  REP(i, N) cnt[i] -= nokori[i] * (nokori[i] + 1) / 2;

  REP(i, N) { cout << N * (N + 1) / 2 + cnt[i] << endl; }

  // dump(kazu)

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<double, double> pdd;
const ll mod = 1e9 + 7;
// const ll mod = 998244353;
#define REP(i, n) for (int i = 0; i < (int)n; ++i)
// debug
#define dump(x) cerr << #x << " = " << (x) << endl;
#define debug(x)                                                               \\
  cerr << #x << " = " << (x) << " (L" << __LINE__ << ")"                       \\
       << " " << __FILE__ << endl;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> v) {
  os << "(" << v.first << ", " << v.second << ")";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << " ";
    }
    os << v[i];
  }
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<vector<T>> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << endl;
    }
    os << v[i];
  }
  return os;
}
string num2bit(ll num, ll len) {
  string bit = "";
  REP(i, len) { bit += char('0' + (num >> i & 1)); }
  return bit;
}

vector<ll> c(202020, 0);
using Graph = vector<vector<int>>;
using pml = pair<map<ll, ll>, ll>;

void printmp(map<ll, ll> &mp) {
  for (auto x : mp) {
    cerr << x << " ";
  }
  cerr << endl;
}

ll N;
ll kazu = 0;
vector<ll> cnt(202020, 0);
pml dfs(ll now, ll par, Graph &G) {
  ll sz = 1;
  map<ll, ll> mp;
  for (auto next : G[now]) {
    if (next == par)
      continue;
    auto [mpch, szch] = dfs(next, now, G);

    sz += szch;
    // ll val = szch - mpch[c[now]];
    ll val = szch;
    if (mpch.find(c[now]) != mpch.end())
      val -= mpch[c[now]];
    cnt[c[now]] -= val * (val + 1) / 2;

    /*
    dump(now)
    cout << "mp: ";
    printmp(mp);
    cout << "mpch: ";
    printmp(mpch);
    cout << endl;
    */

    if ((int)mp.size() < (int)mpch.size())
      swap(mp, mpch);

    /*
    cout << "mp: ";
    printmp(mp);
    cout << "mpch: ";
    printmp(mpch);
    cout << endl;
    */

    for (auto x : mpch) {
      mp[x.first] += x.second;
      kazu++;
    }
    mpch.clear();
  }

  /*
  dump(now)
  printmp(mp);
  dump(sz)
  dump(mp[c[now]])
  */

  // ll val = sz - 1 - mp[c[now]];
  // cnt[c[now]] -= val*(val+1)/2;
  mp[c[now]] = sz;
  pml res = {mp, sz};

  /*
  REP(i, N){
      cout << cnt[i] << " ";
  }
  cout << endl;
  cout << endl;
  */
  return res;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  cin >> N;

  REP(i, N) cin >> c[i], c[i]--;

  Graph G(N);
  REP(i, N - 1) {
    ll a, b;
    cin >> a >> b;
    a--, b--;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  auto [mp, sz] = dfs(0, -1, G);

  // printmp(mp);
  vector<ll> nokori(N, N);

  for (auto x : mp) {
    nokori[x.first] = N - x.second;
  }
  // cout << nokori << endl;
  REP(i, N) cnt[i] -= nokori[i] * (nokori[i] + 1) / 2;

  REP(i, N) { cout << N * (N + 1) / 2 + cnt[i] << endl; }

  // dump(kazu)

  return 0;
}