Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <deque>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stdio.h>
#include <string.h>
#include <string>
#include <vector>
using namespace std;

#define int64 long long
#define pii pair<int64, int64>
#define fr first
#define sc second
#define heap priority_queue
struct piii {
  int64 fr, sc, th;
  bool operator<(const piii &b) {
    if (fr != b.fr)
      return fr < b.fr;
    if (sc != b.sc)
      return sc < b.sc;
    return th < b.th;
  }
};
///////////////////////////////////////////////
///////////////////////////////////////////////

const int maxn = 1e5 + 5;
int n, ccc;
int c[maxn], l[maxn], r[maxn], cnt[maxn], h[maxn], p[maxn];
int lz[maxn * 4], sum[maxn * 4];
vector<int> d[maxn], e[maxn], a;

void input() {
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
    d[c[i]].push_back(i);
  }
  int u, v;
  for (int i = 1; i < n; i++) {
    cin >> u >> v;
    e[u].push_back(v);
    e[v].push_back(u);
  }
}

void dfs(int u) {
  l[u] = ++ccc;
  for (int v : e[u])
    if (v != p[u]) {
      h[v] = h[u] + 1;
      p[v] = u;
      dfs(v);
      cnt[u] += cnt[v];
    }
  cnt[u]++;
  r[u] = ccc;
}

bool cmp(const int &a, const int &b) { return h[a] > h[b]; }

void mdf(int id, int l, int r, int val) {
  if (val == -1)
    return;
  lz[id] = val;
  sum[id] = (r - l + 1) * val;
}

void up(int id, int l, int r, int u, int v, int val) {
  if (v < l || u > r)
    return;
  if (l >= u && r <= v) {
    mdf(id, l, r, val);
    return;
  }
  int mid = (l + r) >> 1;
  mdf(id * 2, l, mid, lz[id]);
  mdf(id * 2 + 1, mid + 1, r, lz[id]);
  lz[id] = -1;
  up(id * 2, l, mid, u, v, val);
  up(id * 2 + 1, mid + 1, r, u, v, val);
  sum[id] = sum[id * 2] + sum[id * 2 + 1];
}

int get(int id, int l, int r, int u, int v) {
  if (v < l || u > r)
    return 0;
  if (l >= u && r <= v)
    return sum[id];
  int mid = (l + r) >> 1;
  mdf(id * 2, l, mid, lz[id]);
  mdf(id * 2 + 1, mid + 1, r, lz[id]);
  return get(id * 2, l, mid, u, v) + get(id * 2 + 1, mid + 1, r, u, v);
}

int main() {
  cout << setprecision(2) << fixed;
  ios::sync_with_stdio(false), cin.tie(0);
  // freopen("input.txt", "r", stdin);
  input();
  dfs(1);
  for (int i = 1; i <= n; i++) {
    sort(d[i].begin(), d[i].end(), cmp);
    up(1, 1, n, 1, n, 1);
    int64 nn = n, res = 0;
    for (int j = 0; j < d[i].size(); j++) {
      int u = d[i][j];
      a.clear();
      int64 tsum = 0;
      for (int v : e[u])
        if (v != p[u]) {
          a.push_back(get(1, 1, n, l[v], r[v]));
          tsum += a[a.size() - 1];
        }
      for (int64 t : a) {
        res += t * (tsum - t);
        tsum -= t;
      }
      int64 tmp = get(1, 1, n, l[d[i][j]], r[d[i][j]]);
      up(1, 1, n, l[d[i][j]], r[d[i][j]], 0);
      nn -= tmp;
      res += tmp * (nn + 1);
    }
    cout << res << endl;
  }
  return 0;
}


Fixed submission:
#include <algorithm>
#include <deque>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stdio.h>
#include <string.h>
#include <string>
#include <vector>
using namespace std;

#define int64 long long
#define pii pair<int64, int64>
#define fr first
#define sc second
#define heap priority_queue
struct piii {
  int64 fr, sc, th;
  bool operator<(const piii &b) {
    if (fr != b.fr)
      return fr < b.fr;
    if (sc != b.sc)
      return sc < b.sc;
    return th < b.th;
  }
};
///////////////////////////////////////////////
///////////////////////////////////////////////

const int maxn = 2e5 + 5;
int n, ccc;
int c[maxn], l[maxn], r[maxn], cnt[maxn], h[maxn], p[maxn];
int lz[maxn * 4], sum[maxn * 4];
vector<int> d[maxn], e[maxn], a;

void input() {
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
    d[c[i]].push_back(i);
  }
  int u, v;
  for (int i = 1; i < n; i++) {
    cin >> u >> v;
    e[u].push_back(v);
    e[v].push_back(u);
  }
}

void dfs(int u) {
  l[u] = ++ccc;
  for (int v : e[u])
    if (v != p[u]) {
      h[v] = h[u] + 1;
      p[v] = u;
      dfs(v);
      cnt[u] += cnt[v];
    }
  cnt[u]++;
  r[u] = ccc;
}

bool cmp(const int &a, const int &b) { return h[a] > h[b]; }

void mdf(int id, int l, int r, int val) {
  if (val == -1)
    return;
  lz[id] = val;
  sum[id] = (r - l + 1) * val;
}

void up(int id, int l, int r, int u, int v, int val) {
  if (v < l || u > r)
    return;
  if (l >= u && r <= v) {
    mdf(id, l, r, val);
    return;
  }
  int mid = (l + r) >> 1;
  mdf(id * 2, l, mid, lz[id]);
  mdf(id * 2 + 1, mid + 1, r, lz[id]);
  lz[id] = -1;
  up(id * 2, l, mid, u, v, val);
  up(id * 2 + 1, mid + 1, r, u, v, val);
  sum[id] = sum[id * 2] + sum[id * 2 + 1];
}

int get(int id, int l, int r, int u, int v) {
  if (v < l || u > r)
    return 0;
  if (l >= u && r <= v)
    return sum[id];
  int mid = (l + r) >> 1;
  mdf(id * 2, l, mid, lz[id]);
  mdf(id * 2 + 1, mid + 1, r, lz[id]);
  return get(id * 2, l, mid, u, v) + get(id * 2 + 1, mid + 1, r, u, v);
}

int main() {
  cout << setprecision(2) << fixed;
  ios::sync_with_stdio(false), cin.tie(0);
  // freopen("input.txt", "r", stdin);
  input();
  dfs(1);
  for (int i = 1; i <= n; i++) {
    sort(d[i].begin(), d[i].end(), cmp);
    up(1, 1, n, 1, n, 1);
    int64 nn = n, res = 0;
    for (int j = 0; j < d[i].size(); j++) {
      int u = d[i][j];
      a.clear();
      int64 tsum = 0;
      for (int v : e[u])
        if (v != p[u]) {
          a.push_back(get(1, 1, n, l[v], r[v]));
          tsum += a[a.size() - 1];
        }
      for (int64 t : a) {
        res += t * (tsum - t);
        tsum -= t;
      }
      int64 tmp = get(1, 1, n, l[d[i][j]], r[d[i][j]]);
      up(1, 1, n, l[d[i][j]], r[d[i][j]], 0);
      nn -= tmp;
      res += tmp * (nn + 1);
    }
    cout << res << endl;
  }
  return 0;
}