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;
typedef vector<vector<ll>> Graph;

const ll mod = 1e9 + 7;
// const ll mod = 998244353;

#define REP(i, n) for (ll i = 0; i < (ll)n; i++)
#define dump(x) cerr << #x << " = " << (x) << endl;
#define spa << " " <<
#define fi first
#define se second

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

template <typename T> void debug(vector<vector<T>> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    cerr << v[i][0];
    for (ll j = 1; j < w; j++)
      cerr spa v[i][j];
    cerr << endl;
  }
};
void debug(vector<string> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    for (ll j = 0; j < w; j++)
      cerr << v[i][j];
    cerr << endl;
  }
};
template <typename T> void debug(vector<T> &v, ll n) {
  if (n != 0)
    cerr << v[0];
  for (ll i = 1; i < n; i++)
    cerr spa v[i];
  cerr << endl;
};

string num2bit(ll num, ll len) {
  string bit = "";
  REP(i, len) { bit += char('0' + (num >> i & 1)); }
  return bit;
}

ll N;
vector<ll> c(202020, -1);
Graph G(202020);
vector<map<ll, ll>> mp(202020);
vector<ll> res(202020, 0);
vector<ll> sz(202020, 0);

ll dfs_sz(ll now, ll par) {
  ll tmp = 1;
  for (auto next : G[now]) {
    if (next == par)
      continue;
    tmp += dfs_sz(next, now);
  }
  sz[now] = tmp;
  return tmp;
}

ll dfs(ll now, ll par) {

  queue<pll> task;
  ll mdx = -1, msz = -1, mnx = -1;
  for (auto next : G[now]) {
    if (next == par)
      continue;
    ll idx = dfs(next, now);
    ll isz = mp[idx].size();
    if (chmax(msz, isz)) {
      mdx = idx;
      mnx = next;
    }
    task.push({idx, next});
  }

  if (mdx == -1) {
    mp[now][c[now]] = 1;
    return now;
  }

  while (!task.empty()) {
    pll tmp = task.front();
    task.pop();
    ll t = tmp.first;
    ll next = tmp.second;
    if (t == mdx)
      continue;

    bool ok = false;
    for (auto x : mp[t]) {
      if (x.first == c[now]) {
        ok = true;
        ll val = sz[next] - mp[t][c[now]];
        res[c[now]] -= val * (val + 1) / 2;
      } else {
        mp[mdx][x.first] += mp[t][x.first];
      }
    }

    if (!ok) {
      ll val = sz[next];
      res[c[now]] -= val * (val + 1) / 2;
    }
  }

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

  mp[mdx][c[now]] = sz[now];

  return mdx;
}

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

  cin >> N;
  REP(i, N) cin >> c[i];
  REP(i, N) c[i]--;
  REP(i, N - 1) {
    ll u, v;
    cin >> u >> v;
    u--, v--;
    G[u].push_back(v);
    G[v].push_back(u);
  }

  REP(i, N) res[i] = N * (N + 1) / 2;

  dfs_sz(0, -1);
  ll mdx = dfs(0, -1);

  REP(i, N) {
    ll val = N - mp[mdx][i];
    res[i] -= val * (val + 1) / 2;
  }

  REP(i, N) { cout << res[i] << endl; }

  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 tmp = dfs(next, now, G);
    auto mpch = tmp.first;
    auto szch = tmp.second;

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