Problem p02710 (C++)

Buggy submission (status: Runtime Error):
// #include "bits/stdc++.h"

#define _USE_MATH_DEFINES
#include <algorithm>
#include <bitset>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

#define rep(i, a, b) for (int i = (a), i##_len = (b); i < i##_len; i++)
#define rrep(i, a, b) for (int i = (b)-1; i >= (a); i--)
#define all(c) begin(c), end(c)

#define int ll
#define SZ(x) ((int)(x).size())
#define pb push_back
#define mp make_pair

// typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<ll, int> pli;
typedef pair<double, double> pdd;
typedef vector<vector<int>> mat;

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 (b < a) {
    a = b;
    return true;
  }
  return false;
}

const int INF =
    sizeof(int) == sizeof(long long) ? 0x3f3f3f3f3f3f3f3fLL : 0x3f3f3f3f;
const int MOD = (int)1e9 + 7;
const double EPS = 1e-9;

int N;
int C[210];
vector<int> G[210];
vector<int> ET;

void dfs(int idx, int par) {
  bool ok = false;
  for (auto e : G[idx]) {
    if (e == par)
      continue;
    ok = true;
    ET.push_back(idx);
    dfs(e, idx);
    ET.push_back(-idx);
  }
  if (!ok) {
    ET.push_back(idx);
    ET.push_back(-idx);
  }
}

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  cin >> N;
  rep(i, 0, N) {
    cin >> C[i + 1];
    C[i + 1]--;
  }

  int a, b;
  rep(i, 0, N - 1) {
    cin >> a >> b;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  dfs(1, -1);

  vector<int> SSZ(SZ(ET));
  set<int> iset;
  vector<map<int, int>> STCM(N);
  vector<map<int, int>> PM(N);
  vector<stack<int>> CS(N);
  vector<int> NGC(N, 0);

  rep(i, 0, N) {
    STCM[i][0] = 0;
    PM[i][0] = 0;
    CS[i].push(0);
  }

  rep(i, 0, SZ(ET)) {
    if (ET[i] > 0) {
      iset.insert(ET[i]);
    }
    SSZ[i] = SZ(iset);

    if (ET[i] > 0) {
      CS[C[ET[i]]].push(ET[i]);
      PM[C[ET[i]]][ET[i]] = i;
      STCM[C[ET[i]]][ET[i]] = 0;
    } else {
      CS[C[-ET[i]]].pop();
      int par = CS[C[-ET[i]]].top();
      int sz = SSZ[i - 1] - SSZ[PM[C[-ET[i]]][-ET[i]]];
      int ngsz = sz - STCM[C[-ET[i]]][-ET[i]];
      NGC[C[-ET[i]]] += ngsz * (ngsz + 1) / 2;
      if (i == SZ(ET) - 1 || ET[i + 1] != -ET[i])
        sz++;
      STCM[C[-ET[i]]][par] += sz;
    }
  }

  rep(i, 0, N) {
    int ngsz = N - STCM[i][0];
    NGC[i] += ngsz * (ngsz + 1) / 2;

    int res = N * (N + 1) / 2 - NGC[i];
    cout << res << endl;
  }

  return 0;
}

Fixed submission:
// #include "bits/stdc++.h"

#define _USE_MATH_DEFINES
#include <algorithm>
#include <bitset>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

#define rep(i, a, b) for (int i = (a), i##_len = (b); i < i##_len; i++)
#define rrep(i, a, b) for (int i = (b)-1; i >= (a); i--)
#define all(c) begin(c), end(c)

#define int ll
#define SZ(x) ((int)(x).size())
#define pb push_back
#define mp make_pair

// typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<ll, int> pli;
typedef pair<double, double> pdd;
typedef vector<vector<int>> mat;

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 (b < a) {
    a = b;
    return true;
  }
  return false;
}

const int INF =
    sizeof(int) == sizeof(long long) ? 0x3f3f3f3f3f3f3f3fLL : 0x3f3f3f3f;
const int MOD = (int)1e9 + 7;
const double EPS = 1e-9;

int N;
int C[200010];
vector<int> G[200010];
vector<int> ET;

void dfs(int idx, int par) {
  bool ok = false;
  for (auto e : G[idx]) {
    if (e == par)
      continue;
    ok = true;
    ET.push_back(idx);
    dfs(e, idx);
    ET.push_back(-idx);
  }
  if (!ok) {
    ET.push_back(idx);
    ET.push_back(-idx);
  }
}

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  cin >> N;
  rep(i, 0, N) {
    cin >> C[i + 1];
    C[i + 1]--;
  }

  int a, b;
  rep(i, 0, N - 1) {
    cin >> a >> b;
    G[a].push_back(b);
    G[b].push_back(a);
  }

  dfs(1, -1);

  vector<int> SSZ(SZ(ET));
  set<int> iset;
  vector<map<int, int>> STCM(N);
  vector<map<int, int>> PM(N);
  vector<stack<int>> CS(N);
  vector<int> NGC(N, 0);

  rep(i, 0, N) {
    STCM[i][0] = 0;
    PM[i][0] = 0;
    CS[i].push(0);
  }

  rep(i, 0, SZ(ET)) {
    if (ET[i] > 0) {
      iset.insert(ET[i]);
    }
    SSZ[i] = SZ(iset);

    if (ET[i] > 0) {
      CS[C[ET[i]]].push(ET[i]);
      PM[C[ET[i]]][ET[i]] = i;
      STCM[C[ET[i]]][ET[i]] = 0;
    } else {
      CS[C[-ET[i]]].pop();
      int par = CS[C[-ET[i]]].top();
      int sz = SSZ[i - 1] - SSZ[PM[C[-ET[i]]][-ET[i]]];
      int ngsz = sz - STCM[C[-ET[i]]][-ET[i]];
      NGC[C[-ET[i]]] += ngsz * (ngsz + 1) / 2;
      if (i == SZ(ET) - 1 || ET[i + 1] != -ET[i])
        sz++;
      STCM[C[-ET[i]]][par] += sz;
    }
  }

  rep(i, 0, N) {
    int ngsz = N - STCM[i][0];
    NGC[i] += ngsz * (ngsz + 1) / 2;

    int res = N * (N + 1) / 2 - NGC[i];
    cout << res << endl;
  }

  return 0;
}