Problem p04008 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <deque>
#include <functional>
#include <iostream>
#include <map>
#include <math.h>
#include <memory.h>
#include <queue>
#include <set>
#include <sstream>
#include <string>
#include <vector>

#pragma comment(linker, "/STACK:64000000")
typedef long long ll;

using namespace std;

const int MAXN = 1 << 18;
const int MOD = 1; // 1000 * 1000 * 1000 + 7;
const int INF = 2e9;
const double PI = acos(-1.0);

int main() {
#ifdef _MSC_VER
  freopen("input.txt", "r", stdin);
#endif

  int n, k;
  while (cin >> n >> k) {
    vector<int> a(n);
    for (int i = 0; i < n; i++)
      cin >> a[i], a[i]--;

    int ans = 0;
    if (a[0] != 0) {
      a[0] = 0;
      ans++;
    }

    vector<int> d(n, -1);
    d[0] = 0;
    function<int(int)> f = [&](int v) {
      if (d[v] != -1)
        return d[v];
      return d[v] = f(a[v]) + 1;
    };
    for (int i = 0; i < n; i++) {
      f(i);
    }
    vector<pair<int, int>> p;
    for (int i = 0; i < n; i++)
      p.push_back(make_pair(-d[i], i));
    sort(p.begin(), p.end());

    vector<int> del(n, 0);
    vector<vector<int>> e(n);
    for (int i = 0; i < n; i++)
      e[a[i]].push_back(i);
    for (int i = 0; i < n; i++) {
      int v = p[i].second;
      if (del[v])
        continue;
      if (-p[i].first <= k)
        continue;
      for (int j = 0; j < k - 1; j++)
        v = a[v];
      a[v] = 0;
      ans++;
      queue<int> q;
      q.push(v);
      while (!q.empty()) {
        int v = q.front();
        q.pop();
        del[v] = 1;
        for (int to : e[v])
          q.push(to);
      }
    }
    cout << ans << endl;
  }

  return 0;
}

Fixed submission:
#include <algorithm>
#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <deque>
#include <functional>
#include <iostream>
#include <map>
#include <math.h>
#include <memory.h>
#include <queue>
#include <set>
#include <sstream>
#include <string>
#include <vector>

#pragma comment(linker, "/STACK:64000000")
typedef long long ll;

using namespace std;

const int MAXN = 1 << 18;
const int MOD = 1; // 1000 * 1000 * 1000 + 7;
const int INF = 2e9;
const double PI = acos(-1.0);

int main() {
#ifdef _MSC_VER
  freopen("input.txt", "r", stdin);
#endif

  int n, k;
  while (cin >> n >> k) {
    vector<int> a(n);
    for (int i = 0; i < n; i++)
      cin >> a[i], a[i]--;

    int ans = 0;
    if (a[0] != 0) {
      a[0] = 0;
      ans++;
    }

    vector<int> d(n, -1);
    d[0] = 0;
    function<int(int)> f = [&](int v) {
      if (d[v] != -1)
        return d[v];
      return d[v] = f(a[v]) + 1;
    };
    for (int i = 0; i < n; i++) {
      f(i);
    }
    vector<pair<int, int>> p;
    for (int i = 0; i < n; i++)
      p.push_back(make_pair(-d[i], i));
    sort(p.begin(), p.end());

    vector<int> del(n, 0);
    vector<vector<int>> e(n);
    for (int i = 0; i < n; i++)
      e[a[i]].push_back(i);
    for (int i = 0; i < n; i++) {
      int v = p[i].second;
      if (del[v])
        continue;
      if (-p[i].first <= k)
        continue;
      for (int j = 0; j < k - 1; j++)
        v = a[v];
      a[v] = 0;
      ans++;
      queue<int> q;
      q.push(v);
      while (!q.empty()) {
        int v = q.front();
        q.pop();
        if (del[v])
          continue;
        del[v] = 1;
        for (int to : e[v])
          q.push(to);
      }
    }
    cout << ans << endl;
  }

  return 0;
}