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

#include <bits/stdc++.h>
#define For(i, a, b) for (int(i) = (int)(a); (i) < (int)(b); ++(i))
#define rFor(i, a, b) for (int(i) = (int)(a)-1; (i) >= (int)(b); --(i))
#define rep(i, n) For((i), 0, (n))
#define rrep(i, n) rFor((i), (n), 0)
#define fi first
#define se second
using namespace std;
typedef long long lint;
typedef unsigned long long ulint;
typedef pair<int, int> pii;
typedef pair<lint, lint> pll;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> T div_floor(T a, T b) {
  if (b < 0)
    a *= -1, b *= -1;
  return a >= 0 ? a / b : (a + 1) / b - 1;
}
template <class T> T div_ceil(T a, T b) {
  if (b < 0)
    a *= -1, b *= -1;
  return a > 0 ? (a - 1) / b + 1 : a / b;
}

constexpr lint mod = 1000000007;
constexpr lint INF = mod * mod;
constexpr int MAX = 200010;

int n, sz[MAX], c[MAX];
vector<int> G[MAX];
map<int, int> mp[MAX];
lint ans[MAX];

int sz_dfs(int v, int pv) {
  sz[v] = 1;
  for (auto nv : G[v])
    if (nv != pv) {
      sz[v] += sz_dfs(nv, v);
    }
  return sz[v];
}

void dfs(int v, int pv) {
  int max_id = -1;
  rep(i, G[v].size()) if (G[v][i] != pv) {
    int nv = G[v][i];
    dfs(nv, v);

    lint tmp = sz[nv] + mp[nv][c[v]];
    ans[c[v]] -= (lint)tmp * (tmp - 1) / 2 + tmp;

    if (max_id < 0 || mp[nv].size() > mp[G[v][max_id]].size()) {
      max_id = i;
    }
  }

  if (max_id >= 0) {
    swap(mp[v], mp[G[v][max_id]]);
    rep(i, G[v].size()) if (i != max_id && G[v][i] != pv) {
      int nv = G[v][i];
      for (auto p : mp[nv])
        mp[v][p.fi] += p.se;
      mp[nv].clear();
    }
  }
  mp[v][c[v]] = -sz[v];
}

int main() {
  scanf("%d", &n);
  rep(i, n) scanf("%d", &c[i]), --c[i];
  rep(i, n - 1) {
    int a, b;
    scanf("%d%d", &a, &b);
    --a;
    --b;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  sz_dfs(0, -1);
  fill(ans, ans + n, (lint)n * (n - 1) / 2 + n);
  dfs(0, -1);
  rep(i, n) {
    lint tmp = n + mp[0][i];
    ans[i] -= (lint)tmp * (tmp - 1) / 2 + tmp;
    printf("%lld\\n", ans[i]);
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define For(i, a, b) for (int(i) = (int)(a); (i) < (int)(b); ++(i))
#define rFor(i, a, b) for (int(i) = (int)(a)-1; (i) >= (int)(b); --(i))
#define rep(i, n) For((i), 0, (n))
#define rrep(i, n) rFor((i), (n), 0)
#define fi first
#define se second
using namespace std;
typedef long long lint;
typedef unsigned long long ulint;
typedef pair<int, int> pii;
typedef pair<lint, lint> pll;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> T div_floor(T a, T b) {
  if (b < 0)
    a *= -1, b *= -1;
  return a >= 0 ? a / b : (a + 1) / b - 1;
}
template <class T> T div_ceil(T a, T b) {
  if (b < 0)
    a *= -1, b *= -1;
  return a > 0 ? (a - 1) / b + 1 : a / b;
}

constexpr lint mod = 1000000007;
constexpr lint INF = mod * mod;
constexpr int MAX = 200010;

int n, sz[MAX], c[MAX];
vector<int> G[MAX];
map<int, int> mp[MAX];
lint ans[MAX];

int sz_dfs(int v, int pv) {
  sz[v] = 1;
  for (auto nv : G[v])
    if (nv != pv) {
      sz[v] += sz_dfs(nv, v);
    }
  return sz[v];
}

void dfs(int v, int pv) {
  int max_id = -1;
  rep(i, G[v].size()) if (G[v][i] != pv) {
    int nv = G[v][i];
    dfs(nv, v);

    lint tmp = sz[nv] + mp[nv][c[v]];
    ans[c[v]] -= (lint)tmp * (tmp - 1) / 2 + tmp;

    if (max_id < 0 || mp[nv].size() > mp[max_id].size()) {
      max_id = i;
    }
  }

  if (max_id >= 0) {
    swap(mp[v], mp[G[v][max_id]]);
    rep(i, G[v].size()) if (i != max_id && G[v][i] != pv) {
      int nv = G[v][i];
      for (auto p : mp[nv])
        mp[v][p.fi] += p.se;
      mp[nv].clear();
    }
  }
  mp[v][c[v]] = -sz[v];
}

int main() {
  scanf("%d", &n);
  rep(i, n) scanf("%d", &c[i]), --c[i];
  rep(i, n - 1) {
    int a, b;
    scanf("%d%d", &a, &b);
    --a;
    --b;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  sz_dfs(0, -1);
  fill(ans, ans + n, (lint)n * (n - 1) / 2 + n);
  dfs(0, -1);
  rep(i, n) {
    lint tmp = n + mp[0][i];
    ans[i] -= (lint)tmp * (tmp - 1) / 2 + tmp;
    printf("%lld\\n", ans[i]);
  }
}