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

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

#define FOR(a, b, c) for (int a = b; a <= c; ++a)
#define FORW(a, b, c) for (int a = b; a >= c; --a)
#define fi first
#define se second
#define pb push_back
#define int long long

typedef pair<int, int> ii;
typedef pair<int, ii> iii;

const int N = 2e5 + 100;
const int maxN = 1e16;
const int oo = 1e18;
const int mod = 1e9 + 7;

int n;
int color_cnt[N], color[N], total[N], sz[N], tmp_count[N];
vector<int> vi[N], group[N];

void prepare(int u, int pu) {
  sz[u] = 1;
  for (auto v : vi[u])
    if (v != pu) {
      prepare(v, u);
      sz[u] += sz[v];
    }
}

void dfs(int u, int pu) {
  int base = tmp_count[color[u]];
  int cur_count = tmp_count[color[u]];
  for (auto v : vi[u])
    if (v != pu) {
      dfs(v, u);
      int inc = tmp_count[color[u]] - cur_count;
      group[color[u]].pb(sz[v] - inc);
      total[color[u]] += (sz[v] - inc);
      cur_count = tmp_count[color[u]];
    }
  tmp_count[color[u]] = base + sz[u];
}

signed main() {
  //    freopen("test.inp", "r", stdin);
  //    freopen("spm.out", "w", stdout);
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> n;
  FOR(i, 1, n) cin >> color[i], color_cnt[color[i]]++;
  FOR(i, 1, n - 1) {
    int u, v;
    cin >> u >> v;
    vi[u].pb(v);
    vi[v].pb(u);
  }
  prepare(1, 1);
  dfs(1, 1);
  FOR(i, 1, n)
  if (total[i] + color_cnt[i] < n) group[i].pb(n - total[i] - color_cnt[i]);

  FOR(i, 1, n) {
    int ans = n * (n - 1) / 2 + n;
    for (auto num : group[i]) {
      ans -= (num * (num - 1) / 2);
      ans -= num;
      // cout << i << ' ' << num << '\\n';
    }
    cout << ans << '\\n';
  }
}


Error: TLE

Fixed submission:

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

#define FOR(a, b, c) for (int a = b; a <= c; ++a)
#define FORW(a, b, c) for (int a = b; a >= c; --a)
#define fi first
#define se second
#define pb push_back
#define int long long

typedef pair<int, int> ii;
typedef pair<int, ii> iii;

const int N = 2e5 + 100;
const int maxN = 1e16;
const int oo = 1e18;
const int mod = 1e9 + 7;

int n, ans, ecnt = 0, tmp_n;
int euler[N], sz[N], remain[N], f[N][21], euler_st[N], euler_en[N], h[N];
bool check[N];
vector<int> vi[N], tmp_vi[N];
vector<int> color[N];

void dfs(int u, int pr) {
  if (u != 1)
    f[u][0] = pr;
  sz[u] = 1;
  euler[u] = euler_st[u] = ++ecnt;
  for (auto v : vi[u])
    if (v != pr) {
      h[v] = h[u] + 1;
      dfs(v, u);
      sz[u] += sz[v];
    }
  euler_en[u] = ecnt;
}

int lca(int u, int v) {
  if (h[u] < h[v])
    swap(u, v);
  FORW(i, 20, 0) if (f[u][i] != -1 && h[f[u][i]] >= h[v]) u = f[u][i];

  FORW(i, 20, 0)
  if (f[u][i] != -1 && f[v][i] != -1 && f[u][i] != f[v][i]) u = f[u][i],
                                                            v = f[v][i];
  if (u != v)
    return f[u][0];
  return u;
}

bool cmp(const int x, const int y) { return euler[x] < euler[y]; }
bool check_in(int u, int pu) {
  return (euler_st[pu] <= euler[u] && euler[u] <= euler_en[pu]);
}

void calc(int u) {
  sort(tmp_vi[u].begin(), tmp_vi[u].end(), cmp);

  for (auto v : tmp_vi[u])
    calc(v);
  int pos = 0, sum = 0, tmp = 0;

  for (auto v : tmp_vi[u]) {
    while (!check_in(v, vi[u][pos]) || vi[u][pos] == f[u][0]) {
      if (vi[u][pos] == f[u][0]) {
        pos += 1;
        continue;
      }
      tmp += sum * sz[vi[u][pos]];
      sum += sz[vi[u][pos]];
      pos += 1;
    }

    tmp += sum * (sz[vi[u][pos]] - sz[v] + remain[v]);
    sum += (sz[vi[u][pos]] - sz[v] + remain[v]);
    pos += 1;
  }

  while (pos < (int32_t)vi[u].size()) {
    if (vi[u][pos] == f[u][0]) {
      pos += 1;
      continue;
    }
    tmp += sum * sz[vi[u][pos]];
    sum += sz[vi[u][pos]];
    pos += 1;
  }

  if (check[u]) {
    remain[u] = 0;
    ans += 1 * tmp + sum + 1; // + u * 2
    ans += (sum + 1) * (tmp_n - sum - 1);
    tmp_n -= (sum + 1);
  } else {
    remain[u] = sum + 1;
  }
}

void solve(int x) {
  vector<int> tmp;
  sort(color[x].begin(), color[x].end(), cmp);

  FOR(i, 0, (int32_t)color[x].size() - 1) {
    check[color[x][i]] = 1;
    tmp.pb(color[x][i]);
    if (i < color[x].size() - 1)
      tmp.pb(lca(color[x][i], color[x][i + 1]));
  }

  sort(tmp.begin(), tmp.end());

  int tmp_size = std::unique(tmp.begin(), tmp.end()) - tmp.begin();
  tmp.resize(tmp_size);
  sort(tmp.begin(), tmp.end(), cmp);

  stack<int> st;
  for (auto v : tmp) {
    while (!st.empty() && lca(st.top(), v) != st.top())
      st.pop();
    if (!st.empty()) {
      tmp_vi[st.top()].pb(v);
    }
    st.push(v);
  }

  tmp_n = n;
  ans = 0;
  if (tmp_size > 0)
    calc(tmp[0]);
  cout << ans << '\\n';

  for (auto v : tmp)
    tmp_vi[v].clear();
  for (auto v : color[x])
    check[v] = 0;
}

signed main() {
  //    freopen("test.inp", "r", stdin);
  //    freopen("spm.out", "w", stdout);
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> n;
  FOR(i, 1, n) {
    int x;
    cin >> x;
    color[x].pb(i);
  }
  FOR(i, 1, n - 1) {
    int x, y;
    cin >> x >> y;
    vi[x].pb(y);
    vi[y].pb(x);
  }

  memset(f, -1, sizeof f);
  dfs(1, 0);

  FOR(i, 1, 19)
  FOR(u, 1, n) if (f[u][i - 1] != -1) f[u][i] = f[f[u][i - 1]][i - 1];
  FOR(i, 1, n) solve(i);
  return 0;
}