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

// #define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define pp pair<int, int>
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define ld long double
#define al(a) (a).begin(), (a).end()
#define mk make_pair
#define check cout << "?" << endl;

ll MOD = 1000000007;
ll mod = 998244353;
int inf = 1000001000;
ll INF = 1e18 + 5;

ll n;
vector<ll> cl, dep, nodes, inodes, visited;
vector<pair<ll, ll>> cpar;
vector<vector<ll>> path, par, seg;
vector<vector<pair<ll, ll>>> fpath;
stack<pair<ll, ll>> st;

void dfs_lca(ll v, ll p, ll d) {
  par[0][v] = p;
  dep[v] = d;
  for (auto to : path[v]) {
    if (to != p)
      dfs_lca(to, v, d + 1);
  }
}

void init(ll n) {
  dfs_lca(0, -1, 0); // #
  rep(k, 19) rep(i, n) {
    if (par[k][i] < 0)
      par[k + 1][i] = -1;
    else
      par[k + 1][i] = par[k][par[k][i]];
  }
}

ll lca(ll u, ll v) {
  if (dep[u] > dep[v])
    swap(u, v);
  rep(k, 20) if ((dep[v] - dep[u]) & (1 << k)) v = par[k][v];
  if (u == v)
    return u;
  for (ll k = 19; k >= 0; k--) {
    if (par[k][u] != par[k][v]) {
      u = par[k][u];
      v = par[k][v];
    }
  }
  return par[0][u];
}

void dfs(ll s) {
  ll cnt = 0;
  while (st.size() > max(dep[s] * 3 - 1, 0LL)) {
    ll x, num;
    tie(x, num) = st.top();
    st.pop();
    if (cnt % 3 != 1)
      cpar[x].second = num;
    else
      cpar[x].first = num;
    cnt++;
  }
  if (cpar[cl[s]].first != -1) {
    ll var = cpar[cl[s]].second;
    if (var == -1)
      var = s;
    fpath[s].push_back(mk(cpar[cl[s]].first, var));
    fpath[cpar[cl[s]].first].push_back(mk(s, var));
    // printf("%lld %lld %lld\\n",s+1,cpar[cl[s]].first+1,var+1);
    visited[s] = 1;
  }
  if (s > 0) {
    ll from = st.top().first;
    st.push(mk(from, -1));
    cpar[from].second = s;
  }
  st.push(mk(cl[s], cpar[cl[s]].first));
  cpar[cl[s]].first = s;
  st.push(mk(cl[s], cpar[cl[s]].second));
  cpar[cl[s]].second = -1;

  for (auto to : path[s]) {
    if (to == par[0][s])
      continue;
    dfs(to);
  }
  if (par[0][s] >= 0) {
    nodes[par[0][s]] += nodes[s];
    inodes[par[0][s]] += inodes[s];
  }
}

void c_dfs(ll s, ll p, ll v) {
  visited[s] = 1;
  for (auto to : fpath[s]) {
    if (to.first == p)
      continue;
    c_dfs(to.first, s, to.second);
  }
  for (auto to : path[s]) {
    if (to == par[0][s])
      continue;
    seg[cl[s]].push_back(inodes[to]);
  }
  if (v != -1) {
    inodes[v] -= nodes[s];
    // printf("%lld %lld\\n",s+1,v+1);
  }
}

int main() {
  cin >> n;
  seg.assign(n, vector<ll>(0));
  path.assign(n, vector<ll>(0));
  fpath.assign(n, vector<pair<ll, ll>>(0));
  par.assign(20, vector<ll>(n, 0));
  dep.assign(n, 0);
  cl.assign(n, 0);
  visited.assign(n, 0);
  cpar.assign(n, mk(-1, -1));
  nodes.assign(n, 1);
  inodes.assign(n, 1);
  vector<ll> ccnt(n, 0);

  rep(i, n) {
    cin >> cl[i];
    cl[i]--;
    ccnt[cl[i]]++;
  }
  rep(i, n - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    path[a].push_back(b);
    path[b].push_back(a);
  }
  init(n);

  dfs(0);
  rep(i, n) {
    if (visited[i] == 0)
      c_dfs(i, -1, -1);
  }

  vector<ll> ans(n, n * (n + 1) / 2);
  rep(i, n) {
    ll node = n - ccnt[i];
    for (auto num : seg[i]) {
      node -= num;
      ans[i] -= num * (num + 1) / 2;
      // printf("%lld ",num);
    }
    ans[i] -= node * (node + 1) / 2;
    // printf("%lld\\n",node);
    printf("%lld\\n", ans[i]);
  }
  // rep(i,n) printf("%lld ",visited[i]);
  // rep(i,n) printf("%lld ",inodes[i]);
  // printf("\\n");
  // rep(i,n) printf("%lld ",nodes[i]);
}

Error: TLE

Fixed submission:

// #define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define pp pair<int, int>
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define ld long double
#define al(a) (a).begin(), (a).end()
#define mk make_pair
#define check cout << "?" << endl;

ll MOD = 1000000007;
ll mod = 998244353;
int inf = 1000001000;
ll INF = 1e18 + 5;

ll n;
vector<ll> cl, dep, cpar, nodes, inodes, visited;
vector<vector<ll>> path, par, seg;
vector<vector<pair<ll, ll>>> fpath;
stack<pair<ll, ll>> st;

void dfs_lca(ll v, ll p, ll d) {
  par[0][v] = p;
  dep[v] = d;
  for (auto to : path[v]) {
    if (to != p)
      dfs_lca(to, v, d + 1);
  }
}

void init(ll n) {
  dfs_lca(0, -1, 0); // #
  rep(k, 19) rep(i, n) {
    if (par[k][i] < 0)
      par[k + 1][i] = -1;
    else
      par[k + 1][i] = par[k][par[k][i]];
  }
}

ll lca(ll u, ll v) {
  if (dep[u] > dep[v])
    swap(u, v);
  rep(k, 20) if ((dep[v] - dep[u]) & (1 << k)) v = par[k][v];
  if (u == v)
    return u;
  for (ll k = 19; k >= 0; k--) {
    if (par[k][u] != par[k][v]) {
      u = par[k][u];
      v = par[k][v];
    }
  }
  return par[0][u];
}

void dfs(ll s) {
  while (st.size() > dep[s]) {
    ll x, num;
    tie(x, num) = st.top();
    st.pop();
    cpar[x] = num;
  }
  if (cpar[cl[s]] != -1) {
    ll var;
    for (auto to : path[cpar[cl[s]]]) {
      if (to == par[0][cpar[cl[s]]])
        continue;
      if (lca(s, to) == to) {
        var = to;
        break;
      }
    }
    fpath[s].push_back(mk(cpar[cl[s]], var));
    fpath[cpar[cl[s]]].push_back(mk(s, var));
    // printf("%lld %lld %lld\\n",s+1,cpar[cl[s]]+1,var+1);
    visited[s] = 1;
  }
  st.push(mk(cl[s], cpar[cl[s]]));
  cpar[cl[s]] = s;

  for (auto to : path[s]) {
    if (to == par[0][s])
      continue;
    dfs(to);
  }
  if (par[0][s] >= 0) {
    nodes[par[0][s]] += nodes[s];
    inodes[par[0][s]] += inodes[s];
  }
}

void c_dfs(ll s, ll p, ll v) {
  visited[s] = 1;
  for (auto to : fpath[s]) {
    if (to.first == p)
      continue;
    c_dfs(to.first, s, to.second);
  }
  for (auto to : path[s]) {
    if (to == par[0][s])
      continue;
    seg[cl[s]].push_back(inodes[to]);
  }
  if (v != -1) {
    inodes[v] -= nodes[s];
    // printf("%lld %lld\\n",s+1,v+1);
  }
}

int main() {
  cin >> n;
  seg.assign(n, vector<ll>(0));
  path.assign(n, vector<ll>(0));
  fpath.assign(n, vector<pair<ll, ll>>(0));
  par.assign(20, vector<ll>(n, 0));
  dep.assign(n, 0);
  cl.assign(n, 0);
  visited.assign(n, 0);
  cpar.assign(n, -1);
  nodes.assign(n, 1);
  inodes.assign(n, 1);
  vector<ll> ccnt(n, 0);

  rep(i, n) {
    cin >> cl[i];
    cl[i]--;
    ccnt[cl[i]]++;
  }
  rep(i, n - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    path[a].push_back(b);
    path[b].push_back(a);
  }
  init(n);

  dfs(0);
  rep(i, n) {
    if (visited[i] == 0)
      c_dfs(i, -1, -1);
  }

  vector<ll> ans(n, n * (n + 1) / 2);
  rep(i, n) {
    ll node = n - ccnt[i];
    for (auto num : seg[i]) {
      node -= num;
      ans[i] -= num * (num + 1) / 2;
      // printf("%lld ",num);
    }
    ans[i] -= node * (node + 1) / 2;
    // printf("%lld\\n",node);
    printf("%lld\\n", ans[i]);
  }
  // rep(i,n) printf("%lld ",visited[i]);
  // rep(i,n) printf("%lld ",inodes[i]);
  // printf("\\n");
  // rep(i,n) printf("%lld ",nodes[i]);
}