Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
using ll = long long;
using P = pair<int, int>;
using vvi = vector<vector<int>>;

std::vector<std::vector<int>> iig(int vn, int en) {
  std::vector<std::vector<int>> res(vn);
  for (int i = 0; i < en; i++) {
    int s, t;
    std::cin >> s >> t;
    s--;
    t--;
    res[s].push_back(t);
    res[t].push_back(s);
  }
  return res;
}

void sep_pc(const std::vector<std::vector<int>> &g, std::vector<int> &p,
            std::vector<std::vector<int>> &c, const int root = 0) {
  int n = g.size();
  std::vector<int> pt(n);
  std::vector<std::vector<int>> ct(n);
  std::function<void(int, int)> f = [&](int x, int parent) {
    pt[x] = parent;
    for (int y : g[x]) {
      if (y == parent)
        continue;
      ct[x].push_back(y);
      f(y, x);
    }
    return;
  };
  f(root, -1);
  p = std::move(pt);
  c = std::move(ct);
  return;
}

void iit(int n, std::vector<int> &p, std::vector<std::vector<int>> &c,
         const int root = 0) {
  sep_pc(iig(n, n - 1), p, c, root);
}

std::vector<int> get_tsize(std::vector<std::vector<int>> &children,
                           int root = 0) {
  std::vector<int> res(children.size());
  std::function<int(int)> f = [&](int x) {
    int a = 1;
    for (int y : children[x]) {
      a += f(y);
    }
    res[x] = a;
    return a;
  };
  f(root);
  return res;
}

int n;
int c[200100];
long long tsize[200100];
vector<int> parent;
vvi children;

vvi categorize_v_by_color() {
  vector<vector<int>> res(n);
  rep(i, n) { res[c[i]].push_back(i); }
  return res;
}

long long *location[200100];
long long head[200100] = {};

void dfs(int x = 0) {
  int color = c[x];
  *location[color] -= tsize[x];
  long long *prev_location = location[color];
  for (int y : children[x]) {
    location[color] = &tsize[y];
    dfs(y);
  }
  location[color] = prev_location;
}

long long comb2(long long x) { return (x + 1) * x / 2; }

int main() {
  cin >> n;
  rep(i, n) {
    int t;
    cin >> t;
    c[i] = t - 1;
  }
  iit(n, parent, children);
  {
    auto tsize_tmp = get_tsize(children, 0);
    rep(i, n) tsize[i] = tsize_tmp[i];
  }
  rep(i, n) head[i] = n;
  rep(i, n) location[i] = head + i;
  dfs();
  auto v_catted = categorize_v_by_color();
  rep(i, n) {
    long long ans = comb2(n) - comb2(head[i]);
    for (int x : v_catted[i]) {
      for (int y : children[x]) {
        ans -= comb2(tsize[y]);
      }
    }
    cout << ans << endl;
  }
}


Error: 0

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
using ll = long long;
using P = pair<int, int>;
using vvi = vector<vector<int>>;

std::vector<std::vector<int>> iig(int vn, int en) {
  std::vector<std::vector<int>> res(vn);
  for (int i = 0; i < en; i++) {
    int s, t;
    std::cin >> s >> t;
    s--;
    t--;
    res[s].push_back(t);
    res[t].push_back(s);
  }
  return res;
}

void sep_pc(const std::vector<std::vector<int>> &g, std::vector<int> &p,
            std::vector<std::vector<int>> &c, const int root = 0) {
  int n = g.size();
  std::vector<int> pt(n);
  std::vector<std::vector<int>> ct(n);
  std::function<void(int, int)> f = [&](int x, int parent) {
    pt[x] = parent;
    for (int y : g[x]) {
      if (y == parent)
        continue;
      ct[x].push_back(y);
      f(y, x);
    }
    return;
  };
  f(root, -1);
  p = std::move(pt);
  c = std::move(ct);
  return;
}

void iit(int n, std::vector<int> &p, std::vector<std::vector<int>> &c,
         const int root = 0) {
  sep_pc(iig(n, n - 1), p, c, root);
}

std::vector<int> get_tsize(std::vector<std::vector<int>> &children,
                           int root = 0) {
  std::vector<int> res(children.size());
  std::function<int(int)> f = [&](int x) {
    int a = 1;
    for (int y : children[x]) {
      a += f(y);
    }
    res[x] = a;
    return a;
  };
  f(root);
  return res;
}

int n;
int c[200010];
long long tsize[20010];
vector<int> parent;
vvi children;

vvi categorize_v_by_color() {
  vector<vector<int>> res(n);
  rep(i, n) { res[c[i]].push_back(i); }
  return res;
}

long long *location[20010];
long long head[20010] = {};

void dfs(int x = 0) {
  int color = c[x];
  *location[color] -= tsize[x];
  long long *prev_location = location[color];
  for (int y : children[x]) {
    location[color] = &tsize[y];
    dfs(y);
  }
  location[color] = prev_location;
}

long long comb2(long long x) { return (x + 1) * x / 2; }

int main() {
  cin >> n;
  rep(i, n) {
    int t;
    cin >> t;
    c[i] = t - 1;
  }
  iit(n, parent, children);
  {
    auto tsize_tmp = get_tsize(children, 0);
    rep(i, n) tsize[i] = tsize_tmp[i];
  }
  rep(i, n) head[i] = n;
  rep(i, n) location[i] = head + i;
  dfs();
  auto v_catted = categorize_v_by_color();
  rep(i, n) {
    long long ans = comb2(n) - comb2(head[i]);
    for (int x : v_catted[i]) {
      for (int y : children[x]) {
        ans -= comb2(tsize[y]);
      }
    }
    cout << ans << endl;
  }
}