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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cctype>
#include <complex>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <tuple>
#include <unistd.h>
#include <vector>
#define _USE_MATH_DEFINES
#define _GLIBCXX_DEBUG
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> plglg;
typedef pair<double, ll> pdlg;
typedef tuple<int, int, int> tiii;
typedef tuple<ll, ll, ll> tlglglg;
typedef tuple<double, double, double> tddd;
typedef complex<double> xy_t;
typedef vector<ll> vll;
typedef vector<vector<ll>> matrix;
typedef unsigned int uint;
typedef unsigned long long ull;
template <class T> using V = vector<T>;
template <class T> using VV = V<V<T>>;
template <class T> using VVV = V<VV<T>>;
#define REP(i, x, y) for (ll i = (ll)x; i < (ll)y; i++)
#define DREP(i, x, y, d) for (ll i = (ll)x; i < (ll)y; i += (ll)d)
#define PER(i, x, y) for (ll i = (ll)x; i > (ll)y; i--)
#define DPER(i, x, y, d) for (ll i = (ll)x; i > (ll)y; i -= (ll)d)
int dx[4] = {1, 0, -1, 0};
int dy[4] = {0, 1, 0, -1};
double pi = 3.141592653589793;
ll mod = 1000000007;
int intmax = 2147483647;
int intmin = -2147483648;
ll llmax = 9223372036854775807;
ll llmin = -9223372036854775807;
int iinf = intmax / 8;
ll inf = llmax / 8;
double eps = 1e-12;
void yesno(bool ok, string yes = "Yes", string no = "No") {
  if (ok) {
    cout << yes << endl;
  } else {
    cout << no << endl;
  }
}

struct edge {
  ll to, cost;
};

int N;
// vector<edge> G[1000000];
VV<edge> G(1000000);
V<ll> d(1000000);

void addedge(VV<edge> &G, int st, int ed, ll co) {
  edge e;
  e.to = ed;
  e.cost = co;
  G[st].push_back(e);
}

// 根付き木を作る。parentVectorに親nodeの番号が入る。
// parentVector[root] = -1
// treeRankVectorにはnodeのrankが入る。rootのrankは0
// subTreeSizeVectorには、そのnodeを根とする部分木のサイズが入る。
vector<ll> par;
vector<ll> treeRank;
vector<ll> subTreeSize;
void makeRootedTree(vector<ll> &parentVector, vector<ll> &treeRankVector,
                    vector<ll> &subTreeSizeVector, ll vecSize, ll root) {
  parentVector = vector<ll>(vecSize);
  treeRankVector = vector<ll>(vecSize);
  assert(root < (ll)parentVector.size());
  parentVector[root] = -1;
  treeRankVector[root] = 0;
  queue<ll> que;
  que.push(root);
  while (!que.empty()) {
    ll node = que.front();
    que.pop();
    REP(i, 0, G[node].size()) {
      edge e = G[node][i];
      if (e.to != parentVector[node]) {
        parentVector[e.to] = node;
        treeRankVector[e.to] = treeRankVector[node] + 1;
        que.push(e.to);
      }
    }
  }

  // ここからはsubTreeSizeVectorの計算
  vector<ll> used(parentVector.size());
  subTreeSizeVector = vector<ll>(vecSize, 0);
  fill(used.begin(), used.end(), 0);
  REP(i, 0, parentVector.size()) {
    if (parentVector[i] != -1 && G[i].size() == 1) {
      que.push(i);
      used[i] = 1;
    }
  }
  while (!que.empty()) {
    ll node = que.front();
    que.pop();
    // rootも含めて子だけのときに何かやりたいときはここに書く
    subTreeSizeVector[node]++;
    if (parentVector[node] == -1) {
      continue;
    }
    ll parent = parentVector[node];
    used[parent]++;
    subTreeSizeVector[parent] += subTreeSizeVector[node];
    // ここに書く
    if ((parentVector[parent] == -1 && used[parent] == (ll)G[parent].size()) ||
        (parentVector[parent] != -1 &&
         used[parent] == (ll)G[parent].size() - 1)) {
      que.push(parent);
    }
  }
}

vll ans(200010, 0);
V<vll> p(200010);
vll c(200010);

void dfs(ll v) {
  // cout << v << " aaa" << endl;
  ll sz = subTreeSize[v];
  // cout << sz << endl;
  ll chinum = sz + sz * (sz - 1) / 2;
  REP(i, 0, G[v].size()) {
    edge e = G[v][i];
    if (e.to == par[v]) {
      continue;
    }
    ll chisz = subTreeSize[e.to];
    chinum -= chisz + chisz * (chisz - 1) / 2;
  }
  ll len = p[c[v]].size();
  ll chiparnum = sz * (p[c[v]][len - 1] - sz);
  // cout << chinum << " " << chiparnum << endl;
  ans[c[v]] += chinum + chiparnum;
  REP(i, 0, G[v].size()) {
    edge e = G[v][i];
    if (e.to == par[v]) {
      continue;
    }
    p[c[v]].push_back(subTreeSize[e.to]);
    dfs(e.to);
  }
  len = p[c[v]].size();
  p[c[v]][len - 1] -= subTreeSize[v];
  if (par[v] != -1) {
    p[c[par[v]]].pop_back();
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  ll N;
  cin >> N;
  REP(i, 0, N) {
    cin >> c[i];
    c[i]--;
  }
  if (N == 1) {
    cout << 1 << endl;
    return 0;
  }
  REP(i, 0, N - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    addedge(G, a, b, 1);
    addedge(G, b, a, 1);
  }
  makeRootedTree(par, treeRank, subTreeSize, N, 0);
  REP(i, 0, N) { p[i].push_back(N); }
  dfs(0);
  // cout << "bbb" << endl;
  REP(i, 0, N) { cout << ans[i] << endl; }
}


Error: 0

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cctype>
#include <complex>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <tuple>
#include <unistd.h>
#include <vector>
#define _USE_MATH_DEFINES
#define _GLIBCXX_DEBUG
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> plglg;
typedef pair<double, ll> pdlg;
typedef tuple<int, int, int> tiii;
typedef tuple<ll, ll, ll> tlglglg;
typedef tuple<double, double, double> tddd;
typedef complex<double> xy_t;
typedef vector<ll> vll;
typedef vector<vector<ll>> matrix;
typedef unsigned int uint;
typedef unsigned long long ull;
template <class T> using V = vector<T>;
template <class T> using VV = V<V<T>>;
template <class T> using VVV = V<VV<T>>;
#define REP(i, x, y) for (ll i = (ll)x; i < (ll)y; i++)
#define DREP(i, x, y, d) for (ll i = (ll)x; i < (ll)y; i += (ll)d)
#define PER(i, x, y) for (ll i = (ll)x; i > (ll)y; i--)
#define DPER(i, x, y, d) for (ll i = (ll)x; i > (ll)y; i -= (ll)d)
int dx[4] = {1, 0, -1, 0};
int dy[4] = {0, 1, 0, -1};
double pi = 3.141592653589793;
ll mod = 1000000007;
int intmax = 2147483647;
int intmin = -2147483648;
ll llmax = 9223372036854775807;
ll llmin = -9223372036854775807;
int iinf = intmax / 8;
ll inf = llmax / 8;
double eps = 1e-12;
void yesno(bool ok, string yes = "Yes", string no = "No") {
  if (ok) {
    cout << yes << endl;
  } else {
    cout << no << endl;
  }
}

struct edge {
  ll to, cost;
};

int N;
// vector<edge> G[1000000];
VV<edge> G(1000000);
V<ll> d(1000000);

void addedge(VV<edge> &G, int st, int ed, ll co) {
  edge e;
  e.to = ed;
  e.cost = co;
  G[st].push_back(e);
}

// 根付き木を作る。parentVectorに親nodeの番号が入る。
// parentVector[root] = -1
// treeRankVectorにはnodeのrankが入る。rootのrankは0
// subTreeSizeVectorには、そのnodeを根とする部分木のサイズが入る。
vector<ll> par;
vector<ll> treeRank;
vector<ll> subTreeSize;
void makeRootedTree(vector<ll> &parentVector, vector<ll> &treeRankVector,
                    vector<ll> &subTreeSizeVector, ll vecSize, ll root) {
  parentVector = vector<ll>(vecSize);
  treeRankVector = vector<ll>(vecSize);
  assert(root < (ll)parentVector.size());
  parentVector[root] = -1;
  treeRankVector[root] = 0;
  queue<ll> que;
  que.push(root);
  while (!que.empty()) {
    ll node = que.front();
    que.pop();
    REP(i, 0, G[node].size()) {
      edge e = G[node][i];
      if (e.to != parentVector[node]) {
        parentVector[e.to] = node;
        treeRankVector[e.to] = treeRankVector[node] + 1;
        que.push(e.to);
      }
    }
  }

  // ここからはsubTreeSizeVectorの計算
  vector<ll> used(parentVector.size());
  subTreeSizeVector = vector<ll>(vecSize, 0);
  fill(used.begin(), used.end(), 0);
  REP(i, 0, parentVector.size()) {
    if (parentVector[i] != -1 && G[i].size() == 1) {
      que.push(i);
      used[i] = 1;
    }
  }
  while (!que.empty()) {
    ll node = que.front();
    que.pop();
    // rootも含めて子だけのときに何かやりたいときはここに書く
    subTreeSizeVector[node]++;
    if (parentVector[node] == -1) {
      continue;
    }
    ll parent = parentVector[node];
    used[parent]++;
    subTreeSizeVector[parent] += subTreeSizeVector[node];
    // ここに書く
    if ((parentVector[parent] == -1 && used[parent] == (ll)G[parent].size()) ||
        (parentVector[parent] != -1 &&
         used[parent] == (ll)G[parent].size() - 1)) {
      que.push(parent);
    }
  }
}

vll ans(200010, 0);
V<vll> p(200010);
vll c(200010);

void dfs(ll v) {
  // cout << v << " aaa" << endl;
  ll sz = subTreeSize[v];
  // cout << sz << endl;
  ll chinum = sz + sz * (sz - 1) / 2;
  REP(i, 0, G[v].size()) {
    edge e = G[v][i];
    if (e.to == par[v]) {
      continue;
    }
    ll chisz = subTreeSize[e.to];
    chinum -= chisz + chisz * (chisz - 1) / 2;
  }
  ll len = p[c[v]].size();
  ll chiparnum = sz * (p[c[v]][len - 1] - sz);
  // cout << chinum << " " << chiparnum << endl;
  ans[c[v]] += chinum + chiparnum;
  REP(i, 0, G[v].size()) {
    edge e = G[v][i];
    if (e.to == par[v]) {
      continue;
    }
    p[c[v]].push_back(subTreeSize[e.to]);
    dfs(e.to);
  }
  p[c[v]].pop_back();
  if (par[v] != -1) {
    len = p[c[par[v]]].size();
    p[c[par[v]]][len - 1] -= subTreeSize[v];
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  ll N;
  cin >> N;
  REP(i, 0, N) {
    cin >> c[i];
    c[i]--;
  }
  if (N == 1) {
    cout << 1 << endl;
    return 0;
  }
  REP(i, 0, N - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    addedge(G, a, b, 1);
    addedge(G, b, a, 1);
  }
  makeRootedTree(par, treeRank, subTreeSize, N, 0);
  REP(i, 0, N) { p[i].push_back(N); }
  dfs(0);
  // cout << "bbb" << endl;
  REP(i, 0, N) { cout << ans[i] << endl; }
}