Problem p04008 (C++)

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

#define int long long
typedef vector<int> vint;
typedef pair<int, int> pint;
typedef vector<pint> vpint;
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, f, n) for (int i = (f); i < (n); i++)
#define all(v) (v).begin(), (v).end()
#define each(it, v)                                                            \
  for (__typeof((v).begin()) it = (v).begin(); it != (v).end(); it++)
#define pb push_back
#define fi first
#define se second
template <typename A, typename B> inline void chmin(A &a, B b) {
  if (a > b)
    a = b;
}
template <typename A, typename B> inline void chmax(A &a, B b) {
  if (a < b)
    a = b;
}

int N, K;
int A[111111];
int dep[111111];
vint G[111111];
void dfs(int v, int p, int d) {
  dep[v] = d;
  for (auto u : G[v]) {
    if (u != p)
      dfs(u, v, d + 1);
  }
}

bool rem[111111];

void col(int v, int p) {
  rem[v] = true;
  for (auto u : G[v]) {
    if (u != p)
      col(u, v);
  }
}
signed main() {
  cin >> N >> K;
  rep(i, N) {
    cin >> A[i];
    A[i]--;
  }

  int cost = 0;
  if (A[0] != 0) {
    A[0] = 0;
    cost++;
  }
  reps(i, 1, N) G[A[i]].pb(i);
  dfs(0, -1, 0);

  vpint vec;
  rep(i, N) vec.pb(pint(dep[i], i));
  sort(all(vec));
  reverse(all(vec));
  for (int i = 0; i < N; i++) {
    int v = vec[i].se;
    int d = vec[i].fi;
    if (rem[v])
      continue;
    if (d <= K)
      continue;
    int u = v;
    for (int j = 0; j < K - 1; j++)
      u = A[u];
    col(u, A[u]);
    cost++;
  }
  cout << cost << endl;
  return 0;
}

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

#define int long long
typedef vector<int> vint;
typedef pair<int, int> pint;
typedef vector<pint> vpint;
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, f, n) for (int i = (f); i < (n); i++)
#define all(v) (v).begin(), (v).end()
#define each(it, v)                                                            \
  for (__typeof((v).begin()) it = (v).begin(); it != (v).end(); it++)
#define pb push_back
#define fi first
#define se second
template <typename A, typename B> inline void chmin(A &a, B b) {
  if (a > b)
    a = b;
}
template <typename A, typename B> inline void chmax(A &a, B b) {
  if (a < b)
    a = b;
}

int N, K;
int A[111111];
int dep[111111];
vint G[111111];
void dfs(int v, int p, int d) {
  dep[v] = d;
  for (auto u : G[v]) {
    if (u != p)
      dfs(u, v, d + 1);
  }
}

bool rem[111111];

void col(int v, int p) {
  if (rem[v])
    return;
  rem[v] = true;
  for (auto u : G[v]) {
    if (u != p)
      col(u, v);
  }
}
signed main() {
  cin >> N >> K;
  rep(i, N) {
    cin >> A[i];
    A[i]--;
  }

  int cost = 0;
  if (A[0] != 0) {
    A[0] = 0;
    cost++;
  }
  reps(i, 1, N) G[A[i]].pb(i);
  dfs(0, -1, 0);

  vpint vec;
  rep(i, N) vec.pb(pint(dep[i], i));
  sort(all(vec));
  reverse(all(vec));
  for (int i = 0; i < N; i++) {
    int v = vec[i].se;
    int d = vec[i].fi;
    if (rem[v])
      continue;
    if (d <= K)
      continue;
    int u = v;
    for (int j = 0; j < K - 1; j++)
      u = A[u];
    col(u, A[u]);
    cost++;
  }
  cout << cost << endl;
  return 0;
}