Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define trav(a, x) for (auto &a : x)
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;

int n;
int c[200005], dead[200005], treesize[200005], major[200005], used[200005];
vi G[200005];
ll sum[200005];
vector<tuple<int, int, ll>> colorsubtrees;

// this DFS finds the centroid.
void dfs1(int n, int p = -1) {
  treesize[n] = 1;
  major[n] = 0;
  for (int h : G[n])
    if (h != p && !dead[h]) {
      dfs1(h, n);
      treesize[n] += treesize[h];
      if (!major[n] || treesize[h] > treesize[major[n]])
        major[n] = h;
    }
}

// this DFS calculates the score.
void dfs2(int n, int p, int treeroot, int centroid) {
  bool uniq = !used[c[n]];
  // printf("DFS2 node = %d (color = %d, color used = %d), treeroot = %d, parent
  // = %d\n", n, c[n], used[c[n]], treeroot, p);
  if (uniq) {
    used[c[n]] = true;
    // printf("Centroid %d, From %d, Subtree %d: %lld paths with color %d\n",
    // centroid, n, treeroot, 1ll * (treesize[centroid] - treesize[treeroot]) *
    // treesize[n], c[n]);
    sum[c[n]] += 1ll * (treesize[centroid] - treesize[treeroot]) * treesize[n];
    // colorsubtrees[c[n]][treeroot] += treesize[n];
    colorsubtrees.push_back({c[n], treeroot, treesize[n]});
  }

  for (int h : G[n])
    if (h != p && !dead[h]) {
      dfs2(h, n, treeroot, centroid);
    }
  if (uniq) {
    used[c[n]] = false;
  }
}

void centroid(int n) {
  dfs1(n);

  int rootsize = treesize[n];
  while (major[n] && treesize[major[n]] * 2 > rootsize) {
    n = major[n];
  }

  // centroid is at n
  // cout << n << " is a centroid" << endl;
  dfs1(n);

  used[c[n]] = true;
  sum[c[n]] += 1; // self -> self
  sum[c[n]] += 1ll * treesize[n] * (treesize[n] - 1) / 2;
  ll paths = 0;
  paths += 1;
  paths += 1ll * treesize[n] * (treesize[n] - 1) / 2;
  // printf("Centroid %d (Color: %d) has paths: %lld + %lld\n", n, c[n], 1ll *
  // treesize[n] * (treesize[n] - 1) / 2, 1ll);
  for (int h : G[n])
    if (!dead[h]) {
      // printf("  .. subtract %lld from subtree %d\n", 1ll * treesize[h] *
      // (treesize[h] - 1) / 2, h);
      sum[c[n]] -= 1ll * treesize[h] * (treesize[h] - 1) / 2;
      paths -= 1ll * treesize[h] * (treesize[h] - 1) / 2;
    }
  // printf("  .. Total: %lld\n", paths);

  // process each subtree
  for (int h : G[n])
    if (!dead[h]) {
      dfs2(h, n, h, n);
    }

  // remove double counts from colorsubtrees
  /*for (auto subtree : colorsubtrees) {
          int color = subtree.first;
          unordered_map<int, ll> sizes = subtree.second;
          ll total = 0;
          ll doubled = 0;

          for (auto p : sizes) total += p.second;
          for (auto p : sizes) {
                  doubled += p.second * (total - p.second);
          }

          if (doubled) {
                  //printf("!! Centroid %d, color %d: doubled %lld\n", n, c[n],
  doubled);
          }
          sum[color] -= doubled / 2;
  }*/
  {
    sort(all(colorsubtrees));

    vector<int> groups;
    groups.push_back(0);
    for (int i = 1; i < colorsubtrees.size(); i++) {
      if (get<0>(colorsubtrees[i]) != get<0>(colorsubtrees[i - 1]))
        groups.push_back(i);
    }
    groups.push_back(colorsubtrees.size());
    for (int g = 0; g < groups.size() - 1; g++) {
      int color = get<0>(colorsubtrees[groups[g]]);
      ll total = 0;
      for (int i = groups[g]; i < groups[g + 1]; i++) {
        total += get<2>(colorsubtrees[i]);
      }

      ll doubled = 0, ctotal = 0;
      for (int i = groups[g]; i < groups[g + 1]; i++) {
        if (i && get<1>(colorsubtrees[i]) != get<1>(colorsubtrees[i - 1])) {
          doubled += ctotal * (total - ctotal);
          ctotal = 0;
        }
        ctotal += get<2>(colorsubtrees[i]);
      }

      doubled += ctotal * (total - ctotal);
      sum[color] -= doubled / 2;
    }
  }
  colorsubtrees.clear();

  used[c[n]] = false;

  // recur
  dead[n] = true;
  for (int h : G[n])
    if (!dead[h])
      centroid(h);
}

int main() {
  cin >> n;
  rep(i, 1, n + 1) cin >> c[i];

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

  centroid(1);
  rep(i, 1, n + 1) { cout << sum[i] << endl; }
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define trav(a, x) for (auto &a : x)
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;

int n;
int c[200005], dead[200005], treesize[200005], major[200005], used[200005];
vi G[200005];
ll sum[200005];
vector<tuple<int, int, ll>> colorsubtrees;

// this DFS finds the centroid.
void dfs1(int n, int p = -1) {
  treesize[n] = 1;
  major[n] = 0;
  for (int h : G[n])
    if (h != p && !dead[h]) {
      dfs1(h, n);
      treesize[n] += treesize[h];
      if (!major[n] || treesize[h] > treesize[major[n]])
        major[n] = h;
    }
}

// this DFS calculates the score.
void dfs2(int n, int p, int treeroot, int centroid) {
  bool uniq = !used[c[n]];
  // printf("DFS2 node = %d (color = %d, color used = %d), treeroot = %d, parent
  // = %d\n", n, c[n], used[c[n]], treeroot, p);
  if (uniq) {
    used[c[n]] = true;
    // printf("Centroid %d, From %d, Subtree %d: %lld paths with color %d\n",
    // centroid, n, treeroot, 1ll * (treesize[centroid] - treesize[treeroot]) *
    // treesize[n], c[n]);
    sum[c[n]] += 1ll * (treesize[centroid] - treesize[treeroot]) * treesize[n];
    // colorsubtrees[c[n]][treeroot] += treesize[n];
    colorsubtrees.push_back({c[n], treeroot, treesize[n]});
  }

  for (int h : G[n])
    if (h != p && !dead[h]) {
      dfs2(h, n, treeroot, centroid);
    }
  if (uniq) {
    used[c[n]] = false;
  }
}

void centroid(int n) {
  dfs1(n);

  int rootsize = treesize[n];
  while (major[n] && treesize[major[n]] * 2 > rootsize) {
    n = major[n];
  }

  // centroid is at n
  // cout << n << " is a centroid" << endl;
  dfs1(n);

  used[c[n]] = true;
  sum[c[n]] += 1; // self -> self
  sum[c[n]] += 1ll * treesize[n] * (treesize[n] - 1) / 2;
  ll paths = 0;
  paths += 1;
  paths += 1ll * treesize[n] * (treesize[n] - 1) / 2;
  // printf("Centroid %d (Color: %d) has paths: %lld + %lld\n", n, c[n], 1ll *
  // treesize[n] * (treesize[n] - 1) / 2, 1ll);
  for (int h : G[n])
    if (!dead[h]) {
      // printf("  .. subtract %lld from subtree %d\n", 1ll * treesize[h] *
      // (treesize[h] - 1) / 2, h);
      sum[c[n]] -= 1ll * treesize[h] * (treesize[h] - 1) / 2;
      paths -= 1ll * treesize[h] * (treesize[h] - 1) / 2;
    }
  // printf("  .. Total: %lld\n", paths);

  // process each subtree
  for (int h : G[n])
    if (!dead[h]) {
      dfs2(h, n, h, n);
    }

  // remove double counts from colorsubtrees
  /*for (auto subtree : colorsubtrees) {
          int color = subtree.first;
          unordered_map<int, ll> sizes = subtree.second;
          ll total = 0;
          ll doubled = 0;

          for (auto p : sizes) total += p.second;
          for (auto p : sizes) {
                  doubled += p.second * (total - p.second);
          }

          if (doubled) {
                  //printf("!! Centroid %d, color %d: doubled %lld\n", n, c[n],
  doubled);
          }
          sum[color] -= doubled / 2;
  }*/
  if (colorsubtrees.size() >= 2) {
    sort(all(colorsubtrees));

    vector<int> groups;
    groups.push_back(0);
    for (int i = 1; i < colorsubtrees.size(); i++) {
      if (get<0>(colorsubtrees[i]) != get<0>(colorsubtrees[i - 1]))
        groups.push_back(i);
    }
    groups.push_back(colorsubtrees.size());
    for (int g = 0; g < groups.size() - 1; g++) {
      int color = get<0>(colorsubtrees[groups[g]]);
      ll total = 0;
      for (int i = groups[g]; i < groups[g + 1]; i++) {
        total += get<2>(colorsubtrees[i]);
      }

      ll doubled = 0, ctotal = 0;
      for (int i = groups[g]; i < groups[g + 1]; i++) {
        if (i && get<1>(colorsubtrees[i]) != get<1>(colorsubtrees[i - 1])) {
          doubled += ctotal * (total - ctotal);
          ctotal = 0;
        }
        ctotal += get<2>(colorsubtrees[i]);
      }

      doubled += ctotal * (total - ctotal);
      sum[color] -= doubled / 2;
    }
  }
  colorsubtrees.clear();

  used[c[n]] = false;

  // recur
  dead[n] = true;
  for (int h : G[n])
    if (!dead[h])
      centroid(h);
}

int main() {
  cin >> n;
  rep(i, 1, n + 1) cin >> c[i];

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

  centroid(1);
  rep(i, 1, n + 1) { cout << sum[i] << endl; }
}