Problem p03499 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 2e5 + 5;
const int MOD = 1e9 + 7;
const int INV2 = MOD - MOD / 2;
typedef long long ll;
struct state {
  int dp[3];
  state(int _a = 0, int _b = 0, int _c = 0) {
    dp[0] = _a, dp[1] = _b, dp[2] = _c;
  }
};
vector<state> s[MAXN];
vector<int> E[MAXN];
int root[MAXN];
ll qpow(ll a, ll b) {
  ll ret = 1;
  while (b) {
    if (b & 1)
      ret = (ret * a) % MOD;
    b >>= 1;
    a = (a * a) % MOD;
  }
  return ret;
}
int merge(int x, int y) {
  if (s[x].size() < s[y].size())
    swap(x, y);
  int x_sz = s[x].size() - 1, y_sz = s[y].size() - 1;
  for (int i = 0; i <= y_sz; i++) {
    int px = x_sz - i, py = y_sz - i;
    int t0, t1, t2;
    t0 = 1LL * s[x][px].dp[0] * s[y][py].dp[0] % MOD;
    t1 = (1LL * s[x][px].dp[1] * s[y][py].dp[0] % MOD +
          1LL * s[x][px].dp[0] * s[y][py].dp[1] % MOD) %
         MOD;
    t2 = (1LL * s[x][px].dp[1] * s[y][py].dp[1] % MOD +
          (1LL * s[x][px].dp[0] + s[x][px].dp[1]) % MOD * s[y][py].dp[2] % MOD +
          (1LL * s[y][py].dp[0] + s[y][py].dp[1]) % MOD * s[x][px].dp[2] % MOD +
          1LL * s[x][px].dp[2] * s[y][py].dp[2] % MOD) %
         MOD;
    s[x][px].dp[0] = t0, s[x][px].dp[1] = t1, s[x][px].dp[2] = t2;
  }
  s[y].clear();
  return x;
}
void dfs(int now) {
  int dep = 0;
  root[now] = now;
  if (E[now].size() == 0) {
    s[root[now]].push_back(state(INV2, INV2, 0));
    return;
  }
  for (int i = 0; i < E[now].size(); i++) {
    int v = E[now][i];
    dfs(v);
    dep = max(dep, (int)s[root[v]].size());
    root[now] = merge(root[now], root[v]);
  }
  int id = root[now];
  int sz = s[id].size() - 1;
  for (int i = 0; i < dep; i++) {
    s[id][sz - i].dp[0] = (s[id][sz - i].dp[0] + s[id][sz - i].dp[2]) % MOD;
    s[id][sz - i].dp[2] = 0;
  }
  s[id].push_back(state(INV2, INV2, 0));
}
inline char nc() {
  static char buf[100000], *p1 = buf, *p2 = buf;
  return p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 100000, stdin), p1 == p2)
             ? EOF
             : *p1++;
}
inline void rea(int &x) {
  char c = nc();
  x = 0;
  for (; c > '9' || c < '0'; c = nc())
    ;
  for (; c >= '0' && c <= '9'; x = x * 10 + c - '0', c = nc())
    ;
}
int main() {
  // freopen("in.txt","r",stdin);
  // freopen("out.txt","w",stdout);
  int n;
  // scanf("%d",&n);
  rea(n);
  for (int i = 1; i <= n; i++) {
    int u;
    // scanf("%d",&u);
    rea(u);
    E[u].push_back(i);
  }
  root[0] = 0;
  dfs(0);
  ll sum = 0;
  for (int i = 0; i < s[root[0]].size(); i++) {
    sum = (sum + s[root[0]][i].dp[1]) % MOD;
  }
  printf("%lld\n", sum * qpow(2, n + 1) % MOD);
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 2e5 + 5;
const int MOD = 1e9 + 7;
const int INV2 = MOD - MOD / 2;
typedef long long ll;
struct state {
  int dp[3];
  state(int _a = 0, int _b = 0, int _c = 0) {
    dp[0] = _a, dp[1] = _b, dp[2] = _c;
  }
};
vector<state> s[MAXN];
vector<int> E[MAXN];
int root[MAXN];
ll qpow(ll a, ll b) {
  ll ret = 1;
  while (b) {
    if (b & 1)
      ret = (ret * a) % MOD;
    b >>= 1;
    a = (a * a) % MOD;
  }
  return ret;
}
int merge(int x, int y) {
  if (s[x].size() < s[y].size())
    swap(x, y);
  int x_sz = s[x].size() - 1, y_sz = s[y].size() - 1;
  for (int i = 0; i <= y_sz; i++) {
    int px = x_sz - i, py = y_sz - i;
    int t0, t1, t2;
    t0 = 1LL * s[x][px].dp[0] * s[y][py].dp[0] % MOD;
    t1 = (1LL * s[x][px].dp[1] * s[y][py].dp[0] % MOD +
          1LL * s[x][px].dp[0] * s[y][py].dp[1] % MOD) %
         MOD;
    t2 = (1LL * s[x][px].dp[1] * s[y][py].dp[1] % MOD +
          (1LL * s[x][px].dp[0] + s[x][px].dp[1]) % MOD * s[y][py].dp[2] % MOD +
          (1LL * s[y][py].dp[0] + s[y][py].dp[1]) % MOD * s[x][px].dp[2] % MOD +
          1LL * s[x][px].dp[2] * s[y][py].dp[2] % MOD) %
         MOD;
    s[x][px].dp[0] = t0, s[x][px].dp[1] = t1, s[x][px].dp[2] = t2;
  }
  s[y].clear();
  return x;
}
void dfs(int now) {
  int dep = 0;
  root[now] = now;
  if (E[now].size() == 0) {
    s[root[now]].push_back(state(INV2, INV2, 0));
    return;
  }
  for (int i = 0; i < E[now].size(); i++) {
    int v = E[now][i];
    dfs(v);
    dep = max(dep, min(int(s[root[now]].size()), (int)s[root[v]].size()));
    root[now] = merge(root[now], root[v]);
  }
  int id = root[now];
  int sz = s[id].size() - 1;
  for (int i = 0; i < dep; i++) {
    s[id][sz - i].dp[0] = (s[id][sz - i].dp[0] + s[id][sz - i].dp[2]) % MOD;
    s[id][sz - i].dp[2] = 0;
  }
  s[id].push_back(state(INV2, INV2, 0));
}
inline char nc() {
  static char buf[100000], *p1 = buf, *p2 = buf;
  return p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 100000, stdin), p1 == p2)
             ? EOF
             : *p1++;
}
inline void rea(int &x) {
  char c = nc();
  x = 0;
  for (; c > '9' || c < '0'; c = nc())
    ;
  for (; c >= '0' && c <= '9'; x = x * 10 + c - '0', c = nc())
    ;
}
int main() {
  // freopen("in.txt","r",stdin);
  // freopen("out.txt","w",stdout);
  int n;
  // scanf("%d",&n);
  rea(n);
  for (int i = 1; i <= n; i++) {
    int u;
    // scanf("%d",&u);
    rea(u);
    E[u].push_back(i);
  }
  root[0] = 0;
  dfs(0);
  ll sum = 0;
  for (int i = 0; i < s[root[0]].size(); i++) {
    sum = (sum + s[root[0]][i].dp[1]) % MOD;
  }
  printf("%lld\n", sum * qpow(2, n + 1) % MOD);
  return 0;
}