Problem p03499 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
const int maxn = 200010;
const int mod = 1e9 + 7;
int n;
vector<int> g[maxn], w[maxn], gg[maxn];
ll dp[maxn][3];
int dep[maxn], size[maxn], fa[maxn][20];
ll ans;
int a[maxn << 1];
stack<int> st;

ll Pow(ll x, ll y) {
  ll ret = 1;
  while (y) {
    if (y & 1)
      ret = ret * x % mod;
    x = x * x % mod;
    y >>= 1;
  }
  return ret;
}

int in[maxn], out[maxn], dfn;
void dfs(int x) {
  in[x] = ++dfn;
  for (int i = 0; i < g[x].size(); i++)
    dfs(g[x][i]);
  out[x] = dfn;
}

bool cmp(int x, int y) { return in[x] < in[y]; }
int lca(int x, int y) {
  if (dep[x] < dep[y])
    swap(x, y);
  for (int i = 19; i >= 0; i--)
    if (dep[fa[x][i]] >= dep[y])
      x = fa[x][i];
  if (x == y)
    return x;
  for (int i = 19; i >= 0; i--)
    if (fa[x][i] != fa[y][i])
      x = fa[x][i], y = fa[y][i];
  return fa[x][0];
}
void solve() {
  size[0] = 1;
  w[0].push_back(0);
  for (int i = 1; i <= n; i++)
    dep[i] = dep[fa[i][0]] + 1, size[dep[i]]++, w[dep[i]].push_back(i);
  for (int d = 0; d <= n; d++) {
    if (!size[d])
      break;
    int tot = 0;
    for (int i = 0; i < w[d].size(); i++)
      a[++tot] = w[d][i];
    sort(a + 1, a + 1 + tot, cmp);
    for (int i = tot; i > 1; i--)
      a[++tot] = lca(a[i], a[i - 1]);
    a[++tot] = 0;
    sort(a + 1, a + 1 + tot, cmp);
    tot = unique(a + 1, a + 1 + tot) - a - 1;
    while (!st.empty())
      st.pop();
    st.push(0);
    for (int i = 2; i <= tot; i++) {
      while (!st.empty() &&
             !(in[st.top()] <= in[a[i]] && out[a[i]] <= out[st.top()]))
        st.pop();
      gg[st.top()].push_back(a[i]);
      st.push(a[i]);
    }
    for (int i = tot; i >= 0; i--) {
      int now = a[i];
      dp[now][0] = dp[now][1] = 0;
      if (dep[now] == d)
        dp[now][0] = dp[now][1] = 1;
      if (dep[now] >= d)
        continue;
      ll dp1 = 0, dp2 = 0, dp0 = 1;
      for (int j = 0; j < gg[now].size(); j++) {
        int v = gg[now][j];
        dp2 = (dp2 * dp[v][0] % mod + (dp2 + dp1) * dp[v][1] % mod) % mod;
        dp1 = (dp1 * dp[v][0] % mod + dp0 * dp[v][1] % mod) % mod;
        dp0 = dp0 * dp[v][0] % mod;
      }
      dp[now][0] = (dp0 + dp2) % mod;
      dp[now][1] = dp1;
    }
    (ans += dp[0][1] * Pow(2, n + 1 - size[d])) %= mod;
    for (int i = 0; i <= n; i++)
      gg[i].clear();
  }
}

int main() {

  scanf("%d", &n);
  for (int i = 1; i <= n; i++)
    scanf("%d", &fa[i][0]), g[fa[i][0]].push_back(i);
  dfs(0);
  for (int j = 1; j <= 19; j++)
    for (int i = 1; i <= n; i++)
      fa[i][j] = fa[fa[i][j - 1]][j - 1];
  solve();
  printf("%lld\n", ans);
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
const int maxn = 200010;
const int mod = 1e9 + 7;
int n;
vector<int> g[maxn], w[maxn], gg[maxn];
ll dp[maxn][3];
int dep[maxn], size[maxn], fa[maxn][20];
ll ans;
int a[maxn << 1];
stack<int> st;

ll Pow(ll x, ll y) {
  ll ret = 1;
  while (y) {
    if (y & 1)
      ret = ret * x % mod;
    x = x * x % mod;
    y >>= 1;
  }
  return ret;
}

int in[maxn], out[maxn], dfn;
void dfs(int x) {
  in[x] = ++dfn;
  for (int i = 0; i < g[x].size(); i++)
    dfs(g[x][i]);
  out[x] = dfn;
}

bool cmp(int x, int y) { return in[x] < in[y]; }
int lca(int x, int y) {
  if (dep[x] < dep[y])
    swap(x, y);
  for (int i = 19; i >= 0; i--)
    if (dep[fa[x][i]] >= dep[y])
      x = fa[x][i];
  if (x == y)
    return x;
  for (int i = 19; i >= 0; i--)
    if (fa[x][i] != fa[y][i])
      x = fa[x][i], y = fa[y][i];
  return fa[x][0];
}
void solve() {
  size[0] = 1;
  w[0].push_back(0);
  for (int i = 1; i <= n; i++)
    dep[i] = dep[fa[i][0]] + 1, size[dep[i]]++, w[dep[i]].push_back(i);
  for (int d = 0; d <= n; d++) {
    if (!size[d])
      break;
    int tot = 0;
    for (int i = 0; i < w[d].size(); i++)
      a[++tot] = w[d][i];
    sort(a + 1, a + 1 + tot, cmp);
    for (int i = tot; i > 1; i--)
      a[++tot] = lca(a[i], a[i - 1]);
    a[++tot] = 0;
    sort(a + 1, a + 1 + tot, cmp);
    tot = unique(a + 1, a + 1 + tot) - a - 1;
    while (!st.empty())
      st.pop();
    st.push(0);
    for (int i = 2; i <= tot; i++) {
      while (!st.empty() &&
             !(in[st.top()] <= in[a[i]] && out[a[i]] <= out[st.top()]))
        st.pop();
      gg[st.top()].push_back(a[i]);
      st.push(a[i]);
    }
    for (int i = tot; i >= 0; i--) {
      int now = a[i];
      dp[now][0] = dp[now][1] = 0;
      if (dep[now] == d)
        dp[now][0] = dp[now][1] = 1;
      if (dep[now] >= d)
        continue;
      ll dp1 = 0, dp2 = 0, dp0 = 1;
      for (int j = 0; j < gg[now].size(); j++) {
        int v = gg[now][j];
        dp2 = (dp2 * dp[v][0] % mod + (dp2 + dp1) * dp[v][1] % mod) % mod;
        dp1 = (dp1 * dp[v][0] % mod + dp0 * dp[v][1] % mod) % mod;
        dp0 = dp0 * dp[v][0] % mod;
      }
      dp[now][0] = (dp0 + dp2) % mod;
      dp[now][1] = dp1;
    }
    (ans += dp[0][1] * Pow(2, n + 1 - size[d])) %= mod;
    for (int i = 1; i <= tot; i++)
      gg[a[i]].clear();
  }
}

int main() {

  scanf("%d", &n);
  for (int i = 1; i <= n; i++)
    scanf("%d", &fa[i][0]), g[fa[i][0]].push_back(i);
  dfs(0);
  for (int j = 1; j <= 19; j++)
    for (int i = 1; i <= n; i++)
      fa[i][j] = fa[fa[i][j - 1]][j - 1];
  solve();
  printf("%lld\n", ans);
  return 0;
}