Problem p03368 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define ri rd<int>
#define rep(i, a, b) for (int i = (a), _ = (b); i <= _; ++i)
#define per(i, a, b) for (int i = (a), _ = (b); i >= _; --i)
#define For(i, a, b) for (int i = (a), _ = (b); i < _; ++i)
const int maxN = 2e5 + 7;
typedef long long LL;

template <class T> T rd() {
  bool f = 1;
  char c = getchar();
  for (; !isdigit(c); c = getchar())
    if (c == '-')
      f = 0;
  T x = 0;
  for (; isdigit(c); c = getchar())
    x = x * 10 + c - '0';
  return f ? x : -x;
}

int n;

struct pii {
  int a[2];
  int id;

  inline bool operator<(const pii &v) const {
    if (a[0] == 0)
      return true;
    if (v.a[0] == 0)
      return false;
    return (LL)a[1] * v.a[0] > (LL)a[0] * v.a[1];
  }

  inline bool operator!=(const pii &v) const {
    return a[0] != v.a[0] || a[1] != v.a[1];
  }
} a[maxN];

int pre[maxN];
int dsu[maxN];
list<int> lis[maxN];
priority_queue<pii> q;

int find(int x) { return dsu[x] != x ? dsu[x] = find(dsu[x]) : x; }

int main() {
#ifndef ONLINE_JUDGE
  freopen("a.in", "r", stdin);
#endif

  n = ri();
  rep(i, 2, n) pre[i] = ri();
  rep(i, 1, n) {
    int x = ri();
    if (x == 1)
      ++a[i].a[1];
    else
      ++a[i].a[0];
    a[i].id = i;
    lis[i].push_back(x);
    q.push(a[i]);
    dsu[i] = i;
  }

  while (!q.empty()) {
    int x = q.top().id;
    if (a[x] != q.top()) {
      q.pop();
      continue;
    }
    q.pop();
    int y = find(pre[x]);
    lis[y].splice(lis[y].end(), lis[x]);
    a[y].a[0] += a[x].a[0];
    a[y].a[1] += a[x].a[1];
    if (y)
      q.push(a[y]);
    dsu[x] = pre[x];
  }

  LL cnt = 0, res = 0;
  for (list<int>::reverse_iterator it = lis[0].rbegin(); it != lis[0].rend();
       ++it) {
    if (*it == 0)
      ++cnt;
    else
      res += cnt;
  }
  printf("%lld\n", res);

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define ri rd<int>
#define rep(i, a, b) for (int i = (a), _ = (b); i <= _; ++i)
#define per(i, a, b) for (int i = (a), _ = (b); i >= _; --i)
#define For(i, a, b) for (int i = (a), _ = (b); i < _; ++i)
const int maxN = 2e5 + 7;
typedef long long LL;

template <class T> T rd() {
  bool f = 1;
  char c = getchar();
  for (; !isdigit(c); c = getchar())
    if (c == '-')
      f = 0;
  T x = 0;
  for (; isdigit(c); c = getchar())
    x = x * 10 + c - '0';
  return f ? x : -x;
}

int n;

struct pii {
  int a[2];
  int id;

  inline bool operator<(const pii &v) const {
    if (a[0] == 0)
      return true;
    if (v.a[0] == 0)
      return false;
    return (LL)a[1] * v.a[0] > (LL)a[0] * v.a[1];
  }

  inline bool operator!=(const pii &v) const {
    return a[0] != v.a[0] || a[1] != v.a[1];
  }
} a[maxN];

int pre[maxN];
int dsu[maxN];
list<int> lis[maxN];
priority_queue<pii> q;

int find(int x) { return dsu[x] != x ? dsu[x] = find(dsu[x]) : x; }

int main() {
#ifndef ONLINE_JUDGE
#endif

  n = ri();
  rep(i, 2, n) pre[i] = ri();
  rep(i, 1, n) {
    int x = ri();
    if (x == 1)
      ++a[i].a[1];
    else
      ++a[i].a[0];
    a[i].id = i;
    lis[i].push_back(x);
    q.push(a[i]);
    dsu[i] = i;
  }

  while (!q.empty()) {
    int x = q.top().id;
    if (a[x] != q.top()) {
      q.pop();
      continue;
    }
    q.pop();
    int y = find(pre[x]);
    lis[y].splice(lis[y].end(), lis[x]);
    a[y].a[0] += a[x].a[0];
    a[y].a[1] += a[x].a[1];
    if (y)
      q.push(a[y]);
    dsu[x] = pre[x];
  }

  LL cnt = 0, res = 0;
  for (list<int>::reverse_iterator it = lis[0].rbegin(); it != lis[0].rend();
       ++it) {
    if (*it == 0)
      ++cnt;
    else
      res += cnt;
  }
  printf("%lld\n", res);

  return 0;
}