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

#include <bits/stdc++.h>
using namespace std;
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
typedef long long ll;
typedef pair<ll, ll> pr;
typedef vector<ll> vc;
typedef unordered_map<ll, ll> umap;
#define pb emplace_back
#define rep(i, a, b) for (ll i = a; i <= b; i++)
#define repr(i, a, b) for (ll i = a; i >= b; i--)
#define reps(i, v) for (ll i = 0; i < v.size(); i++)
template <typename T> void chmin(T &a, const T &b) { a = min(a, b); }
template <typename T> void chmax(T &a, const T &b) { a = max(a, b); }
const ll max_n = 250005;
ll n, c, cr[max_n], full;
ll sz[max_n], depth[max_n], ans[max_n];
ll bg[max_n], bc[max_n];
bool dead[max_n], uc[max_n], us[max_n], ue[max_n];
vector<ll> g[max_n], gt, gs[max_n];

void dfs1(ll now, ll pr) {
  bg[now] = 1;
  for (auto to : g[now])
    if (to != pr && !dead[to]) {
      dfs1(to, now);
      bg[now] += bg[to];
    }
}

void dfs2(ll now, ll pr) {
  gt.pb(now);
  us[cr[now]] = true;
  if (!uc[cr[now]])
    bc[cr[now]] += bg[now];
  else
    ue[now] = true;
  uc[cr[now]] = true;
  for (auto to : g[now])
    if (to != pr && !dead[to]) {
      dfs2(to, now);
    }
  if (!ue[now])
    uc[cr[now]] = false;
  ue[now] = false;
}

void dfs3(ll now, ll pr) {
  gt.pb(now);
  us[cr[now]] = true;
  if (!uc[cr[now]])
    bc[cr[now]] += bg[now];
  else
    ue[now] = true;
  uc[cr[now]] = true;
  for (auto to : g[now])
    if (to != pr && !dead[to]) {
      dfs3(to, now);
    }
  if (cr[now] != cr[c] && !ue[now])
    uc[cr[now]] = false;
  ue[now] = false;
}

int main(void) {
  cin >> n;
  rep(i, 0, n - 1) {
    cin >> cr[i];
    cr[i]--;
  }
  rep(i, 1, n - 1) {
    ll aa, bb;
    cin >> aa >> bb;
    aa--;
    bb--;
    g[aa].pb(bb);
    g[bb].pb(aa);
  }
  auto find_centroid = [&](ll root) {
    auto get_size = [&](auto &&self, ll v, ll p) -> void {
      sz[v] = 1;
      for (auto to : g[v])
        if (to != p && !dead[to]) {
          self(self, to, v);
          sz[v] += sz[to];
        }
    };
    get_size(get_size, root, -1);
    auto dfs = [&](auto &&self, ll v, ll p) -> ll {
      for (auto to : g[v])
        if (to != p && !dead[to]) {
          if (sz[to] > sz[root] / 2)
            return self(self, to, v);
        }
      return v;
    };
    return dfs(dfs, root, -1);
  };
  auto centroid_decomposition = [&](auto &&self, ll root, ll d) -> void {
    c = find_centroid(root);
    dfs1(c, -1);
    // cout << "center " << c << " " << bg[c] << endl;
    dfs2(c, -1);
    bc[cr[c]]--;
    for (auto to : gt) {
      if (us[cr[to]])
        ans[cr[to]] -= bc[cr[to]] * bc[cr[to]];
      // if(us[cr[to]]) cout << " ans1 " << cr[to] << " " << ans[cr[to]] <<
      // endl;
      bc[cr[to]] = 0;
      us[cr[to]] = false;
    }
    gt.clear();
    dead[c] = true;
    depth[c] = d;
    for (auto to : g[c])
      if (!dead[to]) {
        uc[cr[c]] = true;
        us[cr[c]] = true;
        bc[cr[c]] = bg[to];
        gt.pb(c);
        dfs3(to, c);
        for (auto to2 : gt) {
          if (us[cr[to2]])
            ans[cr[to2]] += 2 * bc[cr[to2]] * (bg[c] - bg[to]);
          if (us[cr[to2]])
            ans[cr[to2]] += bc[cr[to2]] * bc[cr[to2]];
          // if(us[cr[to2]]) cout << " ans2 " << cr[to2] << " " << ans[cr[to2]]
          // << endl;
          bc[cr[to2]] = 0;
          us[cr[to2]] = false;
        }
        uc[cr[c]] = false;
        gt.clear();
      }
    ans[cr[c]] += 2;
    // cout << " ans3 " << cr[c] << " " << ans[cr[c]] << endl;
    for (auto to : g[c])
      if (!dead[to]) {
        self(self, to, d + 1);
      }
    dead[c] = false;
  };
  centroid_decomposition(centroid_decomposition, 0, 0);
  rep(i, 0, n - 1) cout << ans[i] / 2 << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> pr;
typedef vector<ll> vc;
#pragma GCC target("avx")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#define pb emplace_back
#define rep(i, a, b) for (ll i = a; i <= b; i++)
#define repr(i, a, b) for (ll i = a; i >= b; i--)
#define reps(i, v) for (ll i = 0; i < v.size(); i++)
struct FastIO {
  FastIO() {
    cin.tie(0);
    ios::sync_with_stdio(0);
  }
} fastiofastio;
#ifdef DEBUG_
#include "../program_contest_library/memo/dump.hpp"
#else
#define dump(...)
#endif
ll n;
bool used[200005] = {0}, used2[200005] = {0};
vector<ll> sz(200005), dead(200005);
vector<ll> sze(200005);
vector<ll> g[200005];
ll dp[200005] = {0}, cs[200005];
vc dm[200005];
ll ans[200005] = {0}, cp;
vc ns, crs;
ll szem;

void dfs2(ll now, ll pr) {
  sze[now] = 1;
  if (!used2[cs[now]]) {
    crs.emplace_back(cs[now]);
    dm[cs[now]].emplace_back(0);
    used2[cs[now]] = true;
  }
  for (auto to : g[now])
    if (to != pr && !dead[to]) {
      dfs2(to, now);
      sze[now] += sze[to];
    }
}

void dfs1(ll now, ll pr) {
  ns.emplace_back(now);
  used[cs[now]] = true;
  for (auto to : g[now])
    if (to != pr && !dead[to]) {
      if (!used[cs[to]])
        dp[cs[to]] += sze[to];
      dfs1(to, now);
    }
  used[cs[now]] = false;
}

int main() {
  // your code goes here
  ll aa, bb, gc;
  cin >> n;
  rep(i, 0, n - 1) {
    cin >> cs[i];
    cs[i]--;
  }
  rep(i, 1, n - 1) {
    cin >> aa >> bb;
    aa--;
    bb--;
    g[aa].pb(bb);
    g[bb].pb(aa);
  }
  auto find_centroid = [&](ll root) {
    auto get_size = [&](auto &&self, ll v, ll p) -> void {
      sz[v] = 1;
      dp[cs[v]] = 0;
      used[cs[v]] = false;
      used2[cs[v]] = false;
      for (auto to : g[v])
        if (to != p && !dead[to]) {
          self(self, to, v);
          sz[v] += sz[to];
        }
    };
    get_size(get_size, root, -1);
    auto dfs = [&](auto &&self, ll v, ll p) -> ll {
      for (auto to : g[v])
        if (to != p && !dead[to]) {
          if (sz[to] > sz[root] / 2)
            return self(self, to, v);
        }
      return v;
    };
    return dfs(dfs, root, -1);
  };
  vector<ll> depth(n);
  szem = n;
  auto centroid_decomposition = [&](auto &&self, ll root, ll d) -> void {
    ll c = find_centroid(root);
    dead[c] = true;
    depth[c] = d;
    cp = cs[c];
    // cout << "center : " << c << " " << szem << " " << cp << endl;
    ll memo[200005] = {0};
    crs.emplace_back(cp);
    used2[cp] = true;
    for (auto to : g[c])
      if (!dead[to]) {
        dfs2(to, c);
        ns.clear();
        dp[cs[to]] += sze[to];
        used[cs[to]] = true;
        dfs1(to, c);
        used[cs[to]] = false;
        ll usc[200005] = {0};
        for (auto in : ns) {
          if (!used[cs[in]])
            dm[cs[in]].emplace_back(dp[cs[in]]);
          if (!used[cs[in]])
            memo[cs[in]] += (dm[cs[in]][dm[cs[in]].size() - 1] -
                             dm[cs[in]][dm[cs[in]].size() - 2]) *
                            (szem - sze[to]);
          used[cs[in]] = true;
        }
        dm[cp].pb(dp[cp]);
        for (auto in : ns) {
          used[cs[in]] = false;
        }
      }
    for (auto ii : crs) {
      if (ii == cp) {
        memo[cp] = 0;
        for (auto to : g[c]) {
          if (dead[to])
            continue;
          ans[cp] += sze[to];
          memo[cp] += sze[to] * (szem - 1 - sze[to]);
        }
        ans[cp] += (memo[cp] / 2);
        // cout << "ans3 : " << ii << " " << ans[ii] << endl;
        dm[cp].clear();
        continue;
      }
      ans[ii] += memo[ii];
      // cout << "ans1 : " << ii << " " << ans[ii] << endl;
      memo[ii] = 0;
      // cout << "dp : " << ii << " " << dp[ii] << endl;
      // cout << ii << " " << sze[ii] << endl;
      memo[ii] -= dp[ii] * dp[ii];
      reps(jj, dm[ii]) {
        if (jj == 0)
          continue;
        memo[ii] +=
            (dm[ii][jj] - dm[ii][jj - 1]) * (dm[ii][jj] - dm[ii][jj - 1]);
      }
      ans[ii] += (memo[ii] / 2);
      dm[ii].clear();
      // cout << "ans2 : " << ii << " " << ans[ii] << endl;
    }
    used[cp] = false;
    crs.clear();
    for (auto to : g[c])
      if (!dead[to]) {
        // dump(root, c, to, d);
        szem = sze[to];
        self(self, to, d + 1);
      }
    dead[c] = false;
  };
  centroid_decomposition(centroid_decomposition, 0, 0);
  rep(i, 0, n - 1) used[i] = false;
  rep(i, 0, n - 1) ans[cs[i]]++;
  rep(i, 0, n - 1) cout << ans[i] << endl;
  return 0;
}