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

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

#define all(x) (x).begin(), (x).end()
#define eb emplace_back
#define mp make_pair
#define pb push_back
#define fi first
#define se second
#define sz(x) ((int)(x).size())
#define rep(x, y, z) for (int x = y; x < z; x++)
#define per(x, z, y) for (int x = z - 1; x >= y; x--)
#define trav(y, x) for (auto &y : x)
#define pii pair<int, int>
#define pll pair<ll, ll>
#define vi vector<int>
#define vpi vector<pii>

int dbg = 0;
#define cri if (0)
#define shame cri cerr << (++dbg) << "\\n\\n";

const int MAX = 2e5 + 1;

int N;
vi col;
vi adj[MAX];
int sbT[MAX];
ll ans[MAX];

ll NC2(ll n) {
  ll x = n - 1;
  return (n * x) / 2;
}

void get(int node, int par) {
  trav(ch, adj[node]) if (ch != par) {
    get(ch, node);
    sbT[node] += sbT[ch];
  }
}
map<int, ll> dfs(int node, int par) {
  ll sub;
  map<int, ll> res;
  trav(ch, adj[node]) {
    if (ch == par)
      continue;
    auto kid = dfs(ch, node);
    auto it = kid.find(col[node]);
    if (it == kid.end())
      sub = 0;
    else
      sub = kid[col[node]];
    ans[col[node]] += NC2(sbT[ch] - sub + 1);

    if (sz(kid) > sz(res))
      swap(res, kid);
    trav(p, kid) {
      it = res.find(p.fi);
      if (it == res.end())
        res[p.fi] = 0;
      res[p.fi] += p.se;
    }
  }
  int c = col[node];
  auto it = res.find(c);
  if (it == res.end())
    res[c] = 0;
  res[c] = sbT[node];
  return res;
}

// aaaaaaaaaaaaaaaa
void solve() {
  scanf("%d", &N);
  int c;
  rep(i, 0, N) {
    scanf("%d", &c);
    col.pb(--c);
  }
  int u, v;
  rep(i, 1, N) {
    scanf("%d%d", &u, &v);
    --u;
    --v;
    adj[u].pb(v);
    adj[v].pb(u);
  }
  fill(sbT, sbT + N, 1);
  get(0, -1);

  fill(ans, ans + N, 0);
  auto rt = dfs(0, -1);
  rep(i, 0, N) {
    auto it = rt.find(i);
    if (it == rt.end())
      rt[i] = 0;
    ll paths = NC2(sbT[0] - rt[i] + 1);
    ans[i] += paths;
  }

  ll tot = NC2(N + 1);
  rep(i, 0, N) {
    ans[i] = tot - ans[i];
    printf("%lld\\n", ans[i]);
  }
}

int main() {
  // int _; scanf("%d", &_); while (_--)
  solve();
  return 0;
}

Error: TLE

Fixed submission:

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

#define all(x) (x).begin(), (x).end()
#define eb emplace_back
#define mp make_pair
#define pb push_back
#define fi first
#define se second
#define sz(x) ((int)(x).size())
#define rep(x, y, z) for (int x = y; x < z; x++)
#define per(x, z, y) for (int x = z - 1; x >= y; x--)
#define trav(y, x) for (auto &y : x)
#define pii pair<int, int>
#define pll pair<ll, ll>
#define vi vector<int>
#define vpi vector<pii>

int dbg = 0;
#define cri if (0)
#define shame cri cerr << (++dbg) << "\\n\\n";

const int MAX = 2e5 + 1;

int N;
vi col;
vi adj[MAX];
int sbT[MAX];
ll ans[MAX];

ll NC2(ll n) {
  ll x = n - 1;
  return (n * x) / 2;
}

void get(int node, int par) {
  trav(ch, adj[node]) if (ch != par) {
    get(ch, node);
    sbT[node] += sbT[ch];
  }
}
map<int, ll> dfs(int node, int par) {
  map<int, ll> res;
  pii big = mp(0, -1);
  vector<map<int, ll>> kids;
  trav(ch, adj[node]) {
    if (ch == par)
      continue;
    int idx = sz(kids);
    kids.pb(dfs(ch, node));
    if (sz(kids[idx]) > big.fi)
      big = mp(sz(kids[idx]), idx);
    ll sub;
    auto it = kids[idx].find(col[node]);
    if (it == kids[idx].end())
      sub = 0;
    else
      sub = kids[idx][col[node]];
    ans[col[node]] += NC2(sbT[ch] - sub + 1);
  }
  if (big.se != -1) {
    res = kids[big.se];
    rep(i, 0, sz(kids)) if (i != big.se) trav(p, kids[i]) {
      auto it = res.find(p.fi);
      if (it == res.end())
        res[p.fi] = 0;
      res[p.fi] += p.se;
    }
  }
  int c = col[node];
  auto it = res.find(c);
  if (it == res.end())
    res[c] = 0;
  res[c] = sbT[node];
  return res;
}

// aaaaaaaaaaaaaaaa
void solve() {
  scanf("%d", &N);
  int c;
  rep(i, 0, N) {
    scanf("%d", &c);
    col.pb(--c);
  }
  int u, v;
  rep(i, 1, N) {
    scanf("%d%d", &u, &v);
    --u;
    --v;
    adj[u].pb(v);
    adj[v].pb(u);
  }
  fill(sbT, sbT + N, 1);
  get(0, -1);

  fill(ans, ans + N, 0);
  auto rt = dfs(0, -1);
  rep(i, 0, N) {
    auto it = rt.find(i);
    if (it == rt.end())
      rt[i] = 0;
    ll paths = NC2(sbT[0] - rt[i] + 1);
    ans[i] += paths;
  }

  ll tot = NC2(N + 1);
  rep(i, 0, N) {
    ans[i] = tot - ans[i];
    printf("%lld\\n", ans[i]);
  }
}

int main() {
  // int _; scanf("%d", &_); while (_--)
  solve();
  return 0;
}