Problem p04008 (C++)

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

struct tree {
  vector<vector<int>> G, son, ancestor;
  vector<int> depth, parent;
  queue<int> que;
  int root = -1;
  bool made_lca = false;

  tree(int n) : G(n), son(n), parent(n, -1), depth(n, -1), ancestor(n) {
    for (int i = 0; i < n; i++)
      ancestor[i].resize(30);
  }

  void add_edge(int u, int v) {
    G[u].push_back(v);
    G[v].push_back(u);
  }

  inline void reset_root() {
    root = -1;
    fill(depth.begin(), depth.end(), -1);
    fill(parent.begin(), parent.end(), -1);
    for (auto &p : ancestor)
      for (int &q : p)
        q = -1;
    for (auto &p : son)
      p.resize(0);
    made_lca = false;
  }

  void have_root(int r) {
    reset_root();
    root = r;
    que.push(r);
    while (que.size()) {
      int p = que.front();
      que.pop();
      for (int q : G[p]) {
        if (parent[p] == q)
          continue;
        son[p].push_back(q);
        parent[q] = p;
        que.push(q);
      }
    }
  }

  void set_depth() {
    assert(~root);
    depth[root] = 0;
    que.push(root);
    while (que.size()) {
      int p = que.front();
      que.pop();
      for (int q : son[p]) {
        depth[q] = depth[p] + 1;
        que.push(q);
      }
    }
  }

  inline void make_lca() {
    assert(~root);
    if (depth[root] < 0)
      set_depth();
    for (int i = 0; i < G.size(); i++)
      ancestor[i][0] = parent[i];
    for (int j = 1; j < 30; j++)
      for (int i = 0; i < G.size(); i++)
        if (~ancestor[i][j - 1])
          ancestor[i][j] = ancestor[ancestor[i][j - 1]][j - 1];
    made_lca = 1;
  }

  int lca(int u, int v) {
    if (!made_lca)
      make_lca();
    if (depth[u] > depth[v])
      swap(u, v);
    int diff = depth[v] - depth[u], cnt = 0;
    while (diff) {
      if (diff & 1)
        v = ancestor[v][cnt];
      cnt++;
      diff >>= 1;
    }
    if (u == v)
      return u;
    for (int k = 28; k >= 0; k--)
      if (ancestor[u][k] != ancestor[v][k])
        u = ancestor[u][k], v = ancestor[v][k];
    return ancestor[u][0];
  }
};
typedef pair<int, int> P;
int used[223456];

signed main() {
  int n, k;
  cin >> n >> k;
  tree t(n);
  int ans = 0;
  int a;
  cin >> a;
  if (a > 1)
    ans++;
  for (int i = 1; i < n; i++) {
    int u;
    cin >> u;
    u--;
    t.add_edge(i, u);
  }
  t.have_root(0);
  t.set_depth();
  priority_queue<P> que;
  queue<int> Q;
  for (int i = 0; i < n; i++)
    if (!t.son[i].size())
      que.push(P(t.depth[i], i));
  while (que.size()) {
    P p = que.top();
    que.pop();
    int idx = p.second;
    if (used[idx]) {
      int tmp = t.parent[idx];
      if (t.depth[tmp] > k)
        que.push(P(t.depth[tmp], tmp));
      continue;
    }
    used[idx] = 1;
    for (int i = 0; i < k - 1; i++) {
      if (!idx)
        break;
      Q.push(t.parent[idx]);
      while (Q.size()) {
        int b = Q.front();
        Q.pop();
        for (int q : t.son[b]) {
          if (used[q])
            continue;
          used[q] = 1;
          Q.push(q);
        }
      }
      idx = t.parent[idx];
      used[idx] = 1;
    }
    if (idx && t.parent[idx]) {
      ans++;
      idx = t.parent[idx];
      if (t.depth[idx] > k)
        que.push(P(t.depth[idx], idx));
    }
  }
  cout << ans << endl;
}


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

struct tree {
  vector<vector<int>> G, son, ancestor;
  vector<int> depth, parent;
  queue<int> que;
  int root = -1;
  bool made_lca = false;

  tree(int n) : G(n), son(n), parent(n, -1), depth(n, -1), ancestor(n) {
    for (int i = 0; i < n; i++)
      ancestor[i].resize(30);
  }

  void add_edge(int u, int v) {
    G[u].push_back(v);
    G[v].push_back(u);
  }

  inline void reset_root() {
    root = -1;
    fill(depth.begin(), depth.end(), -1);
    fill(parent.begin(), parent.end(), -1);
    for (auto &p : ancestor)
      for (int &q : p)
        q = -1;
    for (auto &p : son)
      p.resize(0);
    made_lca = false;
  }

  void have_root(int r) {
    reset_root();
    root = r;
    que.push(r);
    while (que.size()) {
      int p = que.front();
      que.pop();
      for (int q : G[p]) {
        if (parent[p] == q)
          continue;
        son[p].push_back(q);
        parent[q] = p;
        que.push(q);
      }
    }
  }

  void set_depth() {
    assert(~root);
    depth[root] = 0;
    que.push(root);
    while (que.size()) {
      int p = que.front();
      que.pop();
      for (int q : son[p]) {
        depth[q] = depth[p] + 1;
        que.push(q);
      }
    }
  }

  inline void make_lca() {
    assert(~root);
    if (depth[root] < 0)
      set_depth();
    for (int i = 0; i < G.size(); i++)
      ancestor[i][0] = parent[i];
    for (int j = 1; j < 30; j++)
      for (int i = 0; i < G.size(); i++)
        if (~ancestor[i][j - 1])
          ancestor[i][j] = ancestor[ancestor[i][j - 1]][j - 1];
    made_lca = 1;
  }

  int lca(int u, int v) {
    if (!made_lca)
      make_lca();
    if (depth[u] > depth[v])
      swap(u, v);
    int diff = depth[v] - depth[u], cnt = 0;
    while (diff) {
      if (diff & 1)
        v = ancestor[v][cnt];
      cnt++;
      diff >>= 1;
    }
    if (u == v)
      return u;
    for (int k = 28; k >= 0; k--)
      if (ancestor[u][k] != ancestor[v][k])
        u = ancestor[u][k], v = ancestor[v][k];
    return ancestor[u][0];
  }
};
typedef pair<int, int> P;
int used[223456];

signed main() {
  int n, k;
  cin >> n >> k;
  tree t(n);
  int ans = 0;
  int a;
  cin >> a;
  if (a > 1)
    ans++;
  for (int i = 1; i < n; i++) {
    int u;
    cin >> u;
    u--;
    t.add_edge(i, u);
  }
  t.have_root(0);
  t.set_depth();
  priority_queue<P> que;
  queue<int> Q;
  for (int i = 0; i < n; i++)
    if (!t.son[i].size())
      que.push(P(t.depth[i], i));
  while (que.size()) {
    P p = que.top();
    que.pop();
    int idx = p.second;
    if (used[idx]) {
      int tmp = t.parent[idx];
      if (t.depth[tmp] > k && !used[tmp])
        que.push(P(t.depth[tmp], tmp));
      continue;
    }
    used[idx] = 1;
    for (int i = 0; i < k - 1; i++) {
      if (!idx)
        break;
      Q.push(t.parent[idx]);
      while (Q.size()) {
        int b = Q.front();
        Q.pop();
        for (int q : t.son[b]) {
          if (used[q])
            continue;
          used[q] = 1;
          Q.push(q);
        }
      }
      idx = t.parent[idx];
      used[idx] = 1;
    }
    if (idx && t.parent[idx]) {
      ans++;
      idx = t.parent[idx];
      if (t.depth[idx] > k)
        que.push(P(t.depth[idx], idx));
    }
  }
  cout << ans << endl;
}