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

#pragma GCC optimize("Ofast")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx")
#pragma GCC optimize("unroll-loops")

#ifdef _MSC_VER
#pragma warning(disable : 4996)
#include <intrin.h>
#define __builtin_popcount __popcnt
#endif

#include <limits.h>
#include <math.h>
#include <time.h>

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <complex>
#include <cstdio>
#include <cstring>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define REPR(i, n) for (int i = n - 1; i >= 0; --i)
#define FOR(i, m, n) for (int i = m; i < n; ++i)
#define FORR(i, m, n) for (int i = m - 1; i >= n; --i)
#define SORT(v, n) sort(v, v + n);
#define VSORT(v) sort(v.begin(), v.end());
#define REVERSE(v, n) reverse(v, v + n);
#define VREVERSE(v) reverse(v.begin(), v.end())
#define ll long long
#define print(x) cout << (x) << '\\n'
#define pe(x) cout << (x) << " "
#define DEBUG(x) cout << #x << ": " << x << endl
#define lb(v, n) lower_bound(v.begin(), v.end(), (n))
#define ub(v, n) upper_bound(v.begin(), v.end(), (n))
#define int long long
// #define double long double
#define all(x) (x).begin(), (x).end()
#define print_space(v)                                                         \\
  REP(i, v.size()) cout << v[i] << ((i == v.size() - 1) ? "\\n" : " ")
template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
}
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
}
typedef pair<int, int> pii;
typedef array<int, 3> arr3;
std::random_device rd;
std::mt19937 mt(rd());
constexpr ll MOD = 1e9 + 7;
constexpr int MAX = 200020;
const double pi = acos(-1);
constexpr double EPS = 1e-8;
constexpr ll INF = 1e18;

int N;

int col[200020];
vector<int> G[200020];
int sz[200020];
int par[200020];
int calc_size(int n) {
  int ret = 1;
  for (auto nxt : G[n]) {
    if (nxt == par[n])
      continue;
    par[nxt] = n;
    ret += calc_size(nxt);
  }
  return sz[n] = ret;
}
int ans[200020];
void show(map<int, vector<int>> mp) {
  for (auto p : mp) {
    cout << p.first << ": ";
    for (auto v : p.second)
      pe(v);
    cout << endl;
  }
}
void merge(map<int, vector<int>> &mp1, map<int, vector<int>> &mp2) {
  // show(*mp1); show(*mp2);
  if (mp1.size() <= mp2.size()) {
    for (auto itr = mp1.begin(); itr != mp1.end(); itr++) {
      int c = itr->first;
      for (int v : itr->second) {
        mp2[c].push_back(v);
      }
    }
    mp1.swap(mp2);
  } else {
    for (auto itr = mp2.begin(); itr != mp2.end(); itr++) {
      int c = itr->first;
      for (int v : itr->second) {
        mp1[c].push_back(v);
      }
    }
  }
}
void pb(map<int, vector<int>> &mp, int c, int n) { mp[c].push_back(n); }
map<int, vector<int>> dfs(int n) {
  map<int, vector<int>> mp;
  int c = col[n];
  for (auto nxt : G[n]) {
    if (nxt == par[n])
      continue;
    map<int, vector<int>> tmp = move(dfs(nxt));
    int m = sz[nxt];
    // print("a");
    if (tmp.count(c)) {
      for (int &v : tmp[c]) {
        m -= sz[v];
      }
    }
    // cout << "merge" << endl;
    merge(mp, tmp);

    // cout << "merged" << endl;
    // assert(m >= 0);
    ans[c] += m * m - m * (m - 1) / 2;
  }
  mp[c].clear();
  mp[c].push_back(n);
  // print(n+1);
  // for (auto p : mp) {
  //	pe(p.first+1);
  //	cout << ":";
  //	for (auto v : p.second)pe(v+1);
  //	cout << endl;
  // }
  return mp;
}
void solve() {
  cin >> N;
  REP(i, N) {
    cin >> col[i];
    col[i]--;
  }
  REP(i, N - 1) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  calc_size(0);
  map<int, vector<int>> mp;
  mp = dfs(0);
  REP(c, N) {
    int m = sz[0];
    for (int v : mp[c]) {
      m -= sz[v];
    }
    // assert(m >= 0);
    ans[c] += m * m - m * (m - 1) / 2;
    print(N * N - N * (N - 1) / 2 - ans[c]);
  }
}

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  // int q; cin >> q;
  // while (q--)
  solve();
}


Error: TLE

Fixed submission:

#pragma GCC optimize("Ofast")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx")
#pragma GCC optimize("unroll-loops")

#ifdef _MSC_VER
#pragma warning(disable : 4996)
#include <intrin.h>
#define __builtin_popcount __popcnt
#endif

#include <limits.h>
#include <math.h>
#include <time.h>

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <complex>
#include <cstdio>
#include <cstring>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define REPR(i, n) for (int i = n - 1; i >= 0; --i)
#define FOR(i, m, n) for (int i = m; i < n; ++i)
#define FORR(i, m, n) for (int i = m - 1; i >= n; --i)
#define SORT(v, n) sort(v, v + n);
#define VSORT(v) sort(v.begin(), v.end());
#define REVERSE(v, n) reverse(v, v + n);
#define VREVERSE(v) reverse(v.begin(), v.end())
#define ll long long
#define print(x) cout << (x) << '\\n'
#define pe(x) cout << (x) << " "
#define DEBUG(x) cout << #x << ": " << x << endl
#define lb(v, n) lower_bound(v.begin(), v.end(), (n))
#define ub(v, n) upper_bound(v.begin(), v.end(), (n))
#define int long long
// #define double long double
#define all(x) (x).begin(), (x).end()
#define print_space(v)                                                         \\
  REP(i, v.size()) cout << v[i] << ((i == v.size() - 1) ? "\\n" : " ")
template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
}
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
}
typedef pair<int, int> pii;
typedef array<int, 3> arr3;
std::random_device rd;
std::mt19937 mt(rd());
constexpr ll MOD = 1e9 + 7;
constexpr int MAX = 200020;
const double pi = acos(-1);
constexpr double EPS = 1e-8;
constexpr ll INF = 1e18;

int N;

int col[200020];
vector<int> G[200020];
int sz[200020];
int par[200020];
int calc_size(int n) {
  int ret = 1;
  for (auto nxt : G[n]) {
    if (nxt == par[n])
      continue;
    par[nxt] = n;
    ret += calc_size(nxt);
  }
  return sz[n] = ret;
}
int ans[200020];
void show(map<int, vector<int>> mp) {
  for (auto p : mp) {
    cout << p.first << ": ";
    for (auto v : p.second)
      pe(v);
    cout << endl;
  }
}
map<int, vector<int>> merge(map<int, vector<int>> &mp1,
                            map<int, vector<int>> &mp2) {
  // show(*mp1); show(*mp2);
  if (mp1.size() <= mp2.size()) {
    for (auto itr = mp1.begin(); itr != mp1.end(); itr++) {
      int c = itr->first;
      for (int v : itr->second) {
        mp2[c].push_back(v);
      }
    }
    return mp2;
  } else {
    for (auto itr = mp2.begin(); itr != mp2.end(); itr++) {
      int c = itr->first;
      for (int v : itr->second) {
        mp1[c].push_back(v);
      }
    }
    return mp1;
  }
}
void pb(map<int, vector<int>> &mp, int c, int n) { mp[c].push_back(n); }
map<int, vector<int>> dfs(int n) {
  map<int, vector<int>> mp;
  int c = col[n];
  for (auto nxt : G[n]) {
    if (nxt == par[n])
      continue;
    map<int, vector<int>> tmp = move(dfs(nxt));
    int m = sz[nxt];
    // print("a");
    if (tmp.count(c)) {
      for (int &v : tmp[c]) {
        m -= sz[v];
      }
    }
    // cout << "merge" << endl;
    mp = move(merge(mp, tmp));

    // cout << "merged" << endl;
    // assert(m >= 0);
    ans[c] += m * m - m * (m - 1) / 2;
  }
  mp[c].clear();
  mp[c].push_back(n);
  // print(n+1);
  // for (auto p : mp) {
  //	pe(p.first+1);
  //	cout << ":";
  //	for (auto v : p.second)pe(v+1);
  //	cout << endl;
  // }
  return mp;
}
void solve() {
  cin >> N;
  REP(i, N) {
    cin >> col[i];
    col[i]--;
  }
  REP(i, N - 1) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  calc_size(0);
  map<int, vector<int>> mp;
  mp = dfs(0);
  REP(c, N) {
    int m = sz[0];
    for (int v : mp[c]) {
      m -= sz[v];
    }
    // assert(m >= 0);
    ans[c] += m * m - m * (m - 1) / 2;
    print(N * N - N * (N - 1) / 2 - ans[c]);
  }
}

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  // int q; cin >> q;
  // while (q--)
  solve();
}