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

#include <bits/stdc++.h>
using namespace std;
#define int long long
using ll = long long;
using vi = vector<int>;
using pii = pair<int, int>;
#define ALL(c) begin(c), end(c)
#define RALL(c) rbegin(c), rend(c)
#define ITR(i, b, e) for (auto i = (b); i != (e); ++i)
#define FORE(x, c) for (auto &x : c)
#define REPF(i, a, n) for (int i = a, i##len = (int)(n); i < i##len; ++i)
#define REP(i, n) REPF(i, 0, n)
#define REPR(i, n) for (int i = (int)(n); i >= 0; --i)
#define SZ(c) ((int)c.size())
#define CONTAIN(c, x) (c.find(x) != end(c))
#define OUTOFRANGE(y, x, h, w) (y < 0 || x < 0 || y >= h || x >= w)
#define dump(...)
const int DX[9] = {0, 1, 0, -1, 1, 1, -1, -1, 0},
          DY[9] = {-1, 0, 1, 0, -1, 1, 1, -1, 0};
#define INF (1001001001)
#define INFLL (1001001001001001001ll)
template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  ITR(i, begin(v), end(v)) os << *i << (i == end(v) - 1 ? "" : " ");
  return os;
}
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  ITR(i, begin(v), end(v)) is >> *i;
  return is;
}
template <class T> istream &operator>>(istream &is, pair<T, T> &p) {
  is >> p.first >> p.second;
  return is;
}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> using heap = priority_queue<T, vector<T>, greater<T>>;
//------------------------------------------------------------------------------
struct before_main_function {
  before_main_function() {
#ifdef int
#undef INF
#define INF INFLL
#define stoi stoll
#endif
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << setprecision(15) << fixed;
#define endl "\\n"
  }
} before_main_function;
//------------------8<------------------------------------8<--------------------

struct node {
  bool root;
  int child;
  int h;
  node() : root(1), child(-1) {}
};
signed main() {
  int n, k;
  cin >> n >> k;
  vector<node> g(n);
  REP(i, n) {
    int a;
    cin >> a;
    g[i].child = a - 1;
    if (a - 1 > 0)
      g[a - 1].root = 0;
  }
  unordered_set<int> ans;
  REP(i, n) {
    if (g[i].child == -1)
      ans.insert(i);
    else if (g[i].root) {
      int cur = i;
      int d = k;
      while (g[cur].child != -1 && d > 0) {
        ans.insert(cur);
        g[cur].h = d;
        cur = g[cur].child;
        d--;
        if (g[cur].h >= d)
          break;
      }
    }
  }
  cout << SZ(ans) << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define int long long
using ll = long long;
using vi = vector<int>;
using pii = pair<int, int>;
#define ALL(c) begin(c), end(c)
#define RALL(c) rbegin(c), rend(c)
#define ITR(i, b, e) for (auto i = (b); i != (e); ++i)
#define FORE(x, c) for (auto &x : c)
#define REPF(i, a, n) for (int i = a, i##len = (int)(n); i < i##len; ++i)
#define REP(i, n) REPF(i, 0, n)
#define REPR(i, n) for (int i = (int)(n); i >= 0; --i)
#define SZ(c) ((int)c.size())
#define CONTAIN(c, x) (c.find(x) != end(c))
#define OUTOFRANGE(y, x, h, w) (y < 0 || x < 0 || y >= h || x >= w)
#define dump(...)
const int DX[9] = {0, 1, 0, -1, 1, 1, -1, -1, 0},
          DY[9] = {-1, 0, 1, 0, -1, 1, 1, -1, 0};
#define INF (1001001001)
#define INFLL (1001001001001001001ll)
template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  ITR(i, begin(v), end(v)) os << *i << (i == end(v) - 1 ? "" : " ");
  return os;
}
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  ITR(i, begin(v), end(v)) is >> *i;
  return is;
}
template <class T> istream &operator>>(istream &is, pair<T, T> &p) {
  is >> p.first >> p.second;
  return is;
}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> using heap = priority_queue<T, vector<T>, greater<T>>;
//------------------------------------------------------------------------------
struct before_main_function {
  before_main_function() {
#ifdef int
#undef INF
#define INF INFLL
#define stoi stoll
#endif
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << setprecision(15) << fixed;
#define endl "\\n"
  }
} before_main_function;
//------------------8<------------------------------------8<--------------------

struct node {
  bool root;
  int child;
  node() : root(1), child(-1) {}
};
signed main() {
  int n, k;
  cin >> n >> k;
  vector<node> g(n);
  REP(i, n) {
    int a;
    cin >> a;
    g[i].child = a - 1;
    if (a - 1 > 0)
      g[a - 1].root = 0;
  }
  set<int> ans;
  REP(i, n) {
    if (g[i].child == -1)
      ans.insert(i);
    else if (g[i].root) {
      int cur = i;
      int d = 0;
      while (g[cur].child != -1 && d < k) {
        ans.insert(cur);
        cur = g[cur].child;
        d++;
      }
    }
  }
  cout << SZ(ans) << endl;
  return 0;
}