Problem p03368 (C++)

Buggy submission (status: Runtime Error):
#include <cmath>
#include <cstdio>
#include <queue>
#include <tuple>
#include <utility>

using namespace std;

typedef pair<double, int> di;

#define INF 1000000000
#define EPS 1e-9

class UnionFind {
private:
  vector<int> p, ranks, root;

public:
  int initialise(int n) {
    p.assign(n, 1);
    ranks.assign(n, 1);
    root.assign(n, 1);
    for (int i = 0; i <= n; i++)
      p[i] = i, root[i] = i;
  }
  int findSet(int a) {
    if (p[a] == a)
      return a;
    else
      return p[a] = findSet(p[a]);
  }
  bool isSameSet(int a, int b) { return findSet(a) == findSet(b); }
  // a is child b is parent
  void unionSet(int a, int b) {
    if (!isSameSet(a, b)) {
      int pa = findSet(a), pb = findSet(b);
      if (ranks[pa] > ranks[pb]) {
        p[pb] = pa;
        root[pa] = root[pb];
      } else {
        p[pa] = pb;
        if (ranks[pa] == ranks[pb])
          ranks[pb]++;
      }
    }
  }
  int getRoot(int a) { return root[findSet(a)]; }
};

int N;
vector<int> adjList[200005];
int P[200005], V[200005];
priority_queue<di> pq;
int zeroes[200005], ones[200005];
UnionFind ufds;
long long ans = 0;

int main() {
  scanf("%d", &N);
  for (int i = 2; i <= N; i++) {
    scanf("%d", &P[i]);
    adjList[i].push_back(P[i]);
    adjList[P[i]].push_back(i);
  }
  for (int i = 1; i <= N; i++) {
    scanf("%d", &V[i]);
    if (V[i] == 0)
      zeroes[i]++;
    else
      ones[i]++;
  }
  ufds.initialise(N + 5);
  for (int i = 2; i <= N; i++)
    pq.emplace(V[i] == 0 ? INF : 1, i);
  while (!pq.empty()) {
    double r;
    int u;
    tie(r, u) = pq.top();
    pq.pop();
    if (u == 1 || ufds.getRoot(u) != u || (r == INF && ones[u] != 0) ||
        (ones[u] != 0 && fabs(r - (double)zeroes[u] / ones[u]) > EPS))
      continue;
    // printf("%f %d\n", r, u);
    int p = P[u];
    ufds.unionSet(u, p);
    int root = ufds.getRoot(u);
    ans += (long long)ones[root] * zeroes[u];
    zeroes[root] += zeroes[u];
    ones[root] += ones[u];
    pq.emplace(ones[root] == 0 ? INF : (double)zeroes[root] / ones[root], root);
  }
  printf("%lld\n", ans);
  return 0;
}


Fixed submission:
#include <cmath>
#include <cstdio>
#include <queue>
#include <tuple>
#include <utility>

using namespace std;

typedef pair<double, int> di;

#define INF 1000000000
#define EPS 1e-9

class UnionFind {
private:
  vector<int> p, ranks, root;

public:
  int initialise(int n) {
    p.assign(n, 1);
    ranks.assign(n, 1);
    root.assign(n, 1);
    for (int i = 0; i <= n; i++)
      p[i] = i, root[i] = i;
  }
  int findSet(int a) {
    if (p[a] == a)
      return a;
    else
      return p[a] = findSet(p[a]);
  }
  bool isSameSet(int a, int b) { return findSet(a) == findSet(b); }
  // a is child b is parent
  void unionSet(int a, int b) {
    if (!isSameSet(a, b)) {
      int pa = findSet(a), pb = findSet(b);
      if (ranks[pa] > ranks[pb]) {
        p[pb] = pa;
        root[pa] = root[pb];
      } else {
        p[pa] = pb;
        if (ranks[pa] == ranks[pb])
          ranks[pb]++;
      }
    }
  }
  int getRoot(int a) { return root[findSet(a)]; }
};

int N;
vector<int> adjList[200005];
int P[200005], V[200005];
priority_queue<di> pq;
int zeroes[200005], ones[200005];
UnionFind ufds;
long long ans = 0;

int main() {
  scanf("%d", &N);
  if (N == 1) {
    int v;
    scanf("%d", &v);
    ans = 0;
  } else {
    for (int i = 2; i <= N; i++) {
      scanf("%d", &P[i]);
      adjList[i].push_back(P[i]);
      adjList[P[i]].push_back(i);
    }
    for (int i = 1; i <= N; i++) {
      scanf("%d", &V[i]);
      if (V[i] == 0)
        zeroes[i]++;
      else
        ones[i]++;
    }
    ufds.initialise(N + 5);
    for (int i = 2; i <= N; i++)
      pq.emplace(V[i] == 0 ? INF : 1, i);
    while (!pq.empty()) {
      double r;
      int u;
      tie(r, u) = pq.top();
      pq.pop();
      if (u == 1 || ufds.getRoot(u) != u || (r == INF && ones[u] != 0) ||
          (ones[u] != 0 && fabs(r - (double)zeroes[u] / ones[u]) > EPS))
        continue;
      // printf("%f %d\n", r, u);
      int p = P[u];
      ufds.unionSet(u, p);
      int root = ufds.getRoot(u);
      ans += (long long)ones[root] * zeroes[u];
      zeroes[root] += zeroes[u];
      ones[root] += ones[u];
      pq.emplace(ones[root] == 0 ? INF : (double)zeroes[root] / ones[root],
                 root);
    }
  }
  printf("%lld\n", ans);
  return 0;
}