Problem p03368 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
using namespace std;

#define rep(i, n) for (int i = 0; i < (n); i++)
#define all(c) (c).begin(), (c).end()
#define uniq(c) c.erase(unique(all(c)), (c).end())
#define index(xs, x) (int)(lower_bound(all(xs), x) - xs.begin())
#define _1 first
#define _2 second
#define pb push_back
#define INF 1145141919
#define MOD 1000000007
struct UF {
  int U[200000];
  UF(int N) { rep(i, N) U[i] = i; }
  int find(int x) {
    if (U[x] == x)
      return x;
    return U[x] = find(U[x]);
  }
  void unite(int x, int y) {
    x = find(x), y = find(y);
    U[x] = y;
  }
};

int N;
int U[400000];
bool dead[400000];
int A[400000], B[400000]; // 0, 1

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cin >> N;
  U[0] = -1;
  rep(i, N - 1) {
    cin >> U[i + 1];
    U[i + 1]--;
  }
  rep(i, N) {
    int x;
    cin >> x;
    if (x == 0)
      A[i]++;
    else
      B[i]++;
  }
  UF uf(2 * N);

  typedef pair<double, int> P;
  priority_queue<P, vector<P>, greater<P>> q;
  for (int i = 1; i < N; i++)
    q.push(P((double)B[i] / (double)A[i], i));

  int V = N;
  long long sum = 0;
  while (!q.empty()) {
    int x = q.top()._2;
    q.pop();
    if (dead[x])
      continue;
    int p = uf.find(U[x]);
    dead[x] = dead[p] = true;
    int root = V++;
    U[root] = U[p];
    uf.unite(x, root);
    uf.unite(p, root);
    // cout<<x<<"->"<<p<<" (root="<<root<<")\n";
    sum += 1LL * B[p] * A[x];
    A[root] = A[p] + A[x];
    B[root] = B[p] + B[x];
    if (U[root] != -1) {
      q.push(P((double)B[root] / (double)A[root], root));
    }
  }
  cout << sum << "\n";
  return 0;
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
using namespace std;

#define rep(i, n) for (int i = 0; i < (n); i++)
#define all(c) (c).begin(), (c).end()
#define uniq(c) c.erase(unique(all(c)), (c).end())
#define index(xs, x) (int)(lower_bound(all(xs), x) - xs.begin())
#define _1 first
#define _2 second
#define pb push_back
#define INF 1145141919
#define MOD 1000000007
struct UF {
  int U[400000];
  UF(int N) { rep(i, N) U[i] = i; }
  int find(int x) {
    if (U[x] == x)
      return x;
    return U[x] = find(U[x]);
  }
  void unite(int x, int y) {
    x = find(x), y = find(y);
    U[x] = y;
  }
};

int N;
int U[400000];
bool dead[400000];
int A[400000], B[400000]; // 0, 1

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cin >> N;
  U[0] = -1;
  rep(i, N - 1) {
    cin >> U[i + 1];
    U[i + 1]--;
  }
  rep(i, N) {
    int x;
    cin >> x;
    if (x == 0)
      A[i]++;
    else
      B[i]++;
  }
  UF uf(2 * N);

  typedef pair<double, int> P;
  priority_queue<P, vector<P>, greater<P>> q;
  for (int i = 1; i < N; i++)
    q.push(P((double)B[i] / (double)A[i], i));

  int V = N;
  long long sum = 0;
  while (!q.empty()) {
    int x = q.top()._2;
    q.pop();
    if (dead[x])
      continue;
    int p = uf.find(U[x]);
    dead[x] = dead[p] = true;
    int root = V++;
    U[root] = U[p];
    uf.unite(x, root);
    uf.unite(p, root);
    // cout<<x<<"->"<<p<<" (root="<<root<<")\n";
    sum += 1LL * B[p] * A[x];
    A[root] = A[p] + A[x];
    B[root] = B[p] + B[x];
    if (U[root] != -1) {
      q.push(P((double)B[root] / (double)A[root], root));
    }
  }
  cout << sum << "\n";
  return 0;
}