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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <memory>
#include <queue>
#include <random>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

// #pragma comment(linker, "/STACK:16777216")
// #pragma GCC optimize("O3")
// #pragma GCC optimize("Ofast")
// #pragma GCC optimize("unroll-loops")
// #pragma GCC optimize("unswitch-loops")
// #pragma GCC
// target("sse,sse2,sse3,ssse3,sse4,sse4.1,sse4.2,abm,mmx,avx,avx2,popcnt,tune=native")

using namespace std;

mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());

#define ll long long
#define ull unsigned long long
// #define double long double
#define int ll
// #define int ull
// #define int short
// #define int unsigned short
// #define int unsigned int
#define mp make_pair
#define F first
#define S second
#define pii pair<int, int>
#define hash hash228
const int SZ = 2e5 + 4;
const double EPS = 1e-7;
const int mod = 1e9;
const int hashpow = 2;
const int inf = 2e9;

struct custom_hash {
  static uint64_t splitmix64(uint64_t x) {
    // http://xorshift.di.unimi.it/splitmix64.c
    x += 0x9e3779b97f4a7c15;
    x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9;
    x = (x ^ (x >> 27)) * 0x94d049bb133111eb;
    return x ^ (x >> 31);
  }

  size_t operator()(uint64_t x) const {
    static const uint64_t FIXED_RANDOM =
        chrono::steady_clock::now().time_since_epoch().count();
    return splitmix64(x + FIXED_RANDOM);
  }
};

int n;
int c[SZ];
vector<int> g[SZ];

void read() {
  cin >> n;
  for (int i = 0; i < n; ++i) {
    cin >> c[i];
    --c[i];
  }
  for (int i = 1; i < n; ++i) {
    int a, b;
    cin >> a >> b;
    --a, --b;
    g[a].push_back(b);
    g[b].push_back(a);
  }
}

int ans[SZ];
int sz[SZ];
bool used[SZ];

int f(int x) { return (x * x + x) / 2; }

void dfs(int u) {
  used[u] = true;
  sz[u] = 1;
  for (int v : g[u]) {
    if (!used[v]) {
      dfs(v);
      sz[u] += sz[v];
    }
  }
}

vector<pii> tmp[SZ];
int y[SZ];
int z[SZ];
void dfs2(int u) {
  used[u] = true;

  if (!tmp[c[u]].empty())
    y[tmp[c[u]].back().S] -= sz[u];
  else
    z[c[u]] -= sz[u];

  for (int v : g[u]) {
    if (!used[v]) {
      tmp[c[u]].emplace_back(u, v);
      dfs2(v);
      tmp[c[u]].pop_back();
      ans[c[u]] -= f(y[v]);
    }
  }
}

void solve() {
  dfs(0);
  fill(used, used + n, false);
  for (int i = 0; i < n; ++i) {
    ans[i] = f(n);
    y[i] = sz[i];
    z[i] = n;
  }
  dfs2(0);
  for (int i = 0; i < n; ++i) {
    ans[i] -= f(z[i]);
    cout << ans[i] << "\\n";
  }
}

signed main() {

#ifdef __LOCAL
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
#else
  ios_base::sync_with_stdio(0);
  cin.tie(0);
#endif

  //    int t;
  //    cin >> t;
  //    for (int i = 0; i < t; ++i) {
  read();
  solve();
  //    }

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <memory>
#include <queue>
#include <random>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

// #pragma comment(linker, "/STACK:16777216")
#pragma GCC optimize("O3")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("unswitch-loops")
#pragma GCC target(                                                            \\
    "sse,sse2,sse3,ssse3,sse4,sse4.1,sse4.2,abm,mmx,avx,avx2,popcnt,tune=native")

using namespace std;

mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());

#define ll long long
#define ull unsigned long long
// #define double long double
#define int ll
// #define int ull
// #define int short
// #define int unsigned short
// #define int unsigned int
#define mp make_pair
#define F first
#define S second
#define pii pair<int, int>
#define hash hash228
const int SZ = 2e5 + 4;
const double EPS = 1e-7;
const int mod = 1e9;
const int hashpow = 2;
const int inf = 2e9;

struct custom_hash {
  static uint64_t splitmix64(uint64_t x) {
    // http://xorshift.di.unimi.it/splitmix64.c
    x += 0x9e3779b97f4a7c15;
    x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9;
    x = (x ^ (x >> 27)) * 0x94d049bb133111eb;
    return x ^ (x >> 31);
  }

  size_t operator()(uint64_t x) const {
    static const uint64_t FIXED_RANDOM =
        chrono::steady_clock::now().time_since_epoch().count();
    return splitmix64(x + FIXED_RANDOM);
  }
};

int n;
int c[SZ];
vector<int> g[SZ];

void read() {
  cin >> n;
  for (int i = 0; i < n; ++i) {
    cin >> c[i];
    --c[i];
  }
  for (int i = 1; i < n; ++i) {
    int a, b;
    cin >> a >> b;
    --a, --b;
    g[a].push_back(b);
    g[b].push_back(a);
  }
}

int ans[SZ];
int sz[SZ];
bool used[SZ];

int f(int x) { return (x * x + x) / 2; }

vector<int> dfs(int u) {
  used[u] = true;
  sz[u] = 1;
  vector<int> ret(n, 0);
  for (int v : g[u]) {
    if (!used[v]) {
      auto x = dfs(v);
      sz[u] += sz[v];
      ans[c[u]] -= f(sz[v] - x[c[u]]);
      for (int i = 0; i < n; ++i) {
        ret[i] += x[i];
      }
    }
  }
  ret[c[u]] = sz[u];
  return ret;
}

void solve() {
  for (int i = 0; i < n; ++i) {
    ans[i] = f(n);
  }
  auto ret = dfs(0);
  for (int i = 0; i < n; ++i) {
    ans[i] -= f(n - ret[i]);
    cout << ans[i] << "\\n";
  }
}

signed main() {

#ifdef __LOCAL
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
#else
  ios_base::sync_with_stdio(0);
  cin.tie(0);
#endif

  //    int t;
  //    cin >> t;
  //    for (int i = 0; i < t; ++i) {
  read();
  solve();
  //    }

  return 0;
}