Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#define _CRT_SECURE_NO_WARNINGS
#define REL

#pragma GCC target("avx2")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

typedef int64_t i64;
typedef long long LL;
typedef long long ll;
typedef long double LD;
typedef complex<LD> cld;
typedef unsigned int uint;
typedef complex<double> cd;
typedef unsigned long long ULL;

#define endl "\n"
#define dendl "\n\n"
#define mp make_pair
#define mt make_tuple
#define size(v) (int)v.size()
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()

/////////////////////////////////////////////////////////////////////

template <typename T> T sqr(T a) { return a * a; }

template <typename T> int sign(T a) {
  if (a == 0)
    return 0;
  return (a > 0 ? 1 : -1);
}

template <typename T> bool uax(T &a, const T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T> bool uin(T &a, const T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T1, typename T2>
istream &operator>>(istream &in, pair<T1, T2> &p) {
  in >> p.first >> p.second;
  return in;
}

template <typename T1, typename T2>
ostream &operator<<(ostream &out, pair<T1, T2> &p) {
  out << p.first << ' ' << p.second;
  return out;
}

mt19937_64 rnd(0);

const int LOG = 64;
const int N = (int)507;
const int MAXN = (int)3e3 + 7;
const int MOD = (int)998244353;
const int INF = (int)1e9 + 7;
const int CONST = 450;
const LL LINF = (i64)1e17;
const LD PI = 3.1415926535897932384626433832795;
const LD EPS = 1;

/////////////////////////////////////////////////////////////////////

LL ans[N];
bool have[N];
int n, col[N];
vector<int> gr[N];
vector<pair<int, int>> ver[N];
int sz[N], tim, tin[N], tout[N];

bool is_par(int u, int v) { return tin[u] <= tin[v] && tin[v] <= tout[u]; }

void calc_time(int u, int par) {
  tin[u] = tim++;
  for (int v : gr[u]) {
    if (v == par)
      continue;
    calc_time(v, u);
  }
  tout[u] = tim++;
}

void dfs(int u, int par) {
  sz[u] = 1;
  for (int v : gr[u]) {
    if (v == par)
      continue;
    dfs(v, u);
    sz[u] += sz[v];
  }
  int pos = col[u];
  reverse(all(gr[u]));
  for (int v : gr[u]) {
    if (v == par)
      continue;
    int comp = sz[v];
    while (size(ver[pos])) {
      int w, cnt;
      tie(w, cnt) = ver[pos].back();
      if (!is_par(v, w))
        break;
      comp -= cnt;
      ver[pos].pop_back();
    }
    ans[pos - 1] -= 1LL * comp * (comp + 1) / 2;
  }
  ver[pos].push_back(mp(u, sz[u]));
  if (u == 0) {
    for (int i = 1; i <= n; i++) {
      if (i == col[u])
        continue;
      int comp = n;
      for (auto it : ver[i])
        comp -= it.second;
      ans[i - 1] -= 1LL * comp * (comp + 1) / 2;
    }
  }
}

void solve() {
  cin >> n;
  for (int i = 0; i < n; i++) {
    cin >> col[i];
    have[col[i] - 1] = true;
    ans[i] = 1LL * n * (n + 1) / 2;
  }
  for (int i = 1; i < n; i++) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    gr[u].push_back(v);
    gr[v].push_back(u);
  }
  calc_time(0, 0);
  dfs(0, 0);
  for (int i = 0; i < n; i++) {
    cout << (have[i] ? ans[i] : 0) << endl;
  }
}

int main() {
  double start = clock();
#ifdef _DEBUG
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
  cout << fixed << setprecision(20);
  cerr << fixed << setprecision(20);
#else
  cout << fixed << setprecision(20);
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  rnd.seed(time(0));
#endif

  int t = 1;
  for (int i = 0; i < t; i++) {
    solve();
  }

#ifdef _DEBUG
  cout << endl << "Time is " << (clock() - start) / 1000.0;
#endif
}

Fixed submission:
#define _CRT_SECURE_NO_WARNINGS
#define REL

#pragma GCC target("avx2")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

typedef int64_t i64;
typedef long long LL;
typedef long long ll;
typedef long double LD;
typedef complex<LD> cld;
typedef unsigned int uint;
typedef complex<double> cd;
typedef unsigned long long ULL;

#define endl "\n"
#define dendl "\n\n"
#define mp make_pair
#define mt make_tuple
#define size(v) (int)v.size()
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()

/////////////////////////////////////////////////////////////////////

template <typename T> T sqr(T a) { return a * a; }

template <typename T> int sign(T a) {
  if (a == 0)
    return 0;
  return (a > 0 ? 1 : -1);
}

template <typename T> bool uax(T &a, const T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T> bool uin(T &a, const T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T1, typename T2>
istream &operator>>(istream &in, pair<T1, T2> &p) {
  in >> p.first >> p.second;
  return in;
}

template <typename T1, typename T2>
ostream &operator<<(ostream &out, pair<T1, T2> &p) {
  out << p.first << ' ' << p.second;
  return out;
}

mt19937_64 rnd(0);

const int LOG = 64;
const int N = (int)2e5 + 7;
const int MAXN = (int)3e3 + 7;
const int MOD = (int)998244353;
const int INF = (int)1e9 + 7;
const int CONST = 450;
const LL LINF = (i64)1e17;
const LD PI = 3.1415926535897932384626433832795;
const LD EPS = 1;

/////////////////////////////////////////////////////////////////////

LL ans[N];
bool have[N];
int n, col[N];
vector<int> gr[N];
vector<pair<int, int>> ver[N];
int sz[N], tim, tin[N], tout[N];

bool is_par(int u, int v) { return tin[u] <= tin[v] && tin[v] <= tout[u]; }

void calc_time(int u, int par) {
  tin[u] = tim++;
  for (int v : gr[u]) {
    if (v == par)
      continue;
    calc_time(v, u);
  }
  tout[u] = tim++;
}

void dfs(int u, int par) {
  sz[u] = 1;
  for (int v : gr[u]) {
    if (v == par)
      continue;
    dfs(v, u);
    sz[u] += sz[v];
  }
  int pos = col[u];
  reverse(all(gr[u]));
  for (int v : gr[u]) {
    if (v == par)
      continue;
    int comp = sz[v];
    while (size(ver[pos])) {
      int w, cnt;
      tie(w, cnt) = ver[pos].back();
      if (!is_par(v, w))
        break;
      comp -= cnt;
      ver[pos].pop_back();
    }
    ans[pos - 1] -= 1LL * comp * (comp + 1) / 2;
  }
  ver[pos].push_back(mp(u, sz[u]));
  if (u == 0) {
    for (int i = 1; i <= n; i++) {
      if (i == col[u])
        continue;
      int comp = n;
      for (auto it : ver[i])
        comp -= it.second;
      ans[i - 1] -= 1LL * comp * (comp + 1) / 2;
    }
  }
}

void solve() {
  cin >> n;
  for (int i = 0; i < n; i++) {
    cin >> col[i];
    have[col[i] - 1] = true;
    ans[i] = 1LL * n * (n + 1) / 2;
  }
  for (int i = 1; i < n; i++) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    gr[u].push_back(v);
    gr[v].push_back(u);
  }
  calc_time(0, 0);
  dfs(0, 0);
  for (int i = 0; i < n; i++) {
    cout << (have[i] ? ans[i] : 0) << endl;
  }
}

int main() {
  double start = clock();
#ifdef _DEBUG
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
  cout << fixed << setprecision(20);
  cerr << fixed << setprecision(20);
#else
  cout << fixed << setprecision(20);
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  rnd.seed(time(0));
#endif

  int t = 1;
  for (int i = 0; i < t; i++) {
    solve();
  }

#ifdef _DEBUG
  cout << endl << "Time is " << (clock() - start) / 1000.0;
#endif
}