Problem p01933 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

#define _overload(_1, _2, _3, name, ...) name
#define _rep(i, n) _range(i, 0, n)
#define _range(i, a, b) for (int i = int(a); i < int(b); ++i)
#define rep(...) _overload(__VA_ARGS__, _range, _rep, )(__VA_ARGS__)

#define _rrep(i, n) _rrange(i, n, 0)
#define _rrange(i, a, b) for (int i = int(a) - 1; i >= int(b); --i)
#define rrep(...) _overload(__VA_ARGS__, _rrange, _rrep, )(__VA_ARGS__)

#define _all(arg) begin(arg), end(arg)
#define uniq(arg) sort(_all(arg)), (arg).erase(unique(_all(arg)), end(arg))
#define getidx(ary, key) lower_bound(_all(ary), key) - begin(ary)
#define clr(a, b) memset((a), (b), sizeof(a))
#define bit(n) (1LL << (n))
#define popcount(n) (__builtin_popcountll(n))

using namespace std;

template <class T> bool chmax(T &a, const T &b) {
  return (a < b) ? (a = b, 1) : 0;
}
template <class T> bool chmin(T &a, const T &b) {
  return (b < a) ? (a = b, 1) : 0;
}

using ll = long long;
using R = long double;
const R EPS = 1e-9L; // [-1000,1000]->EPS=1e-8 [-10000,10000]->EPS=1e-7
inline int sgn(const R &r) { return (r > EPS) - (r < -EPS); }
inline R sq(R x) { return sqrt(max(x, 0.0L)); }

const int dx[8] = {1, 0, -1, 0, 1, -1, -1, 1};
const int dy[8] = {0, 1, 0, -1, 1, 1, -1, -1};

// Problem Specific Parameter:

const int limit = 100010;
vector<int> graph[limit];
set<int> ans;
int n, k;

int dfs(int v, int p, int d) {

  if (graph[v].size() == 0) {
    ans.insert(v);
    return d;
  }

  int ret = 1 << 20;
  for (auto &v2 : graph[v]) {
    if (v2 == p)
      continue;
    ret = min(ret, dfs(v2, v, d + 1));
  }

  if (ret - d < k)
    ans.insert(v);

  return ret;
}

int a[limit];

int main(void) {
  cin >> n >> k;

  rep(i, n) {
    cin >> a[i];
    a[i]--;
    if (a[i] >= 0)
      graph[a[i]].push_back(i);
  }

  rep(i, n) {
    if (a[i] == -1) {
      dfs(i, -1, 0);
      ans.insert(i);
    }
  }

  cout << ans.size() << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

#define _overload(_1, _2, _3, name, ...) name
#define _rep(i, n) _range(i, 0, n)
#define _range(i, a, b) for (int i = int(a); i < int(b); ++i)
#define rep(...) _overload(__VA_ARGS__, _range, _rep, )(__VA_ARGS__)

#define _rrep(i, n) _rrange(i, n, 0)
#define _rrange(i, a, b) for (int i = int(a) - 1; i >= int(b); --i)
#define rrep(...) _overload(__VA_ARGS__, _rrange, _rrep, )(__VA_ARGS__)

#define _all(arg) begin(arg), end(arg)
#define uniq(arg) sort(_all(arg)), (arg).erase(unique(_all(arg)), end(arg))
#define getidx(ary, key) lower_bound(_all(ary), key) - begin(ary)
#define clr(a, b) memset((a), (b), sizeof(a))
#define bit(n) (1LL << (n))
#define popcount(n) (__builtin_popcountll(n))

using namespace std;

template <class T> bool chmax(T &a, const T &b) {
  return (a < b) ? (a = b, 1) : 0;
}
template <class T> bool chmin(T &a, const T &b) {
  return (b < a) ? (a = b, 1) : 0;
}

using ll = long long;
using R = long double;
const R EPS = 1e-9L; // [-1000,1000]->EPS=1e-8 [-10000,10000]->EPS=1e-7
inline int sgn(const R &r) { return (r > EPS) - (r < -EPS); }
inline R sq(R x) { return sqrt(max(x, 0.0L)); }

const int dx[8] = {1, 0, -1, 0, 1, -1, -1, 1};
const int dy[8] = {0, 1, 0, -1, 1, 1, -1, -1};

// Problem Specific Parameter:

set<int> ans;
int a[100010];
int d[100010];

int main(void) {
  int n, k;
  cin >> n >> k;

  rep(i, n) {
    cin >> a[i];
    a[i]--;
    if (a[i] >= 0)
      d[a[i]]++;
  }

  rep(i, n) {
    if (a[i] == -1) {
      ans.insert(i);
    } else if (d[i] == 0) {

      int cur = i;
      rep(loop, k) {
        ans.insert(cur);
        if (a[cur] == -1)
          break;
        cur = a[cur];
      }
    }
  }

  cout << ans.size() << endl;

  return 0;
}