Problem p04008 (C++)

Buggy submission (status: Time Limit Exceeded):
/*
 */

#pragma comment(linker, "/STACK:16777216")
#define _CRT_SECURE_NO_WARNINGS

#include <algorithm>
#include <assert.h>
#include <complex>
#include <ctime>
#include <fstream>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <memory.h>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <vector>

#define y0 sdkfaslhagaklsldk
#define y1 aasdfasdfasdf
#define yn askfhwqriuperikldjk
#define j1 assdgsdgasghsf
#define tm sdfjahlfasfh
#define lr asgasgash
#define norm asdfasdgasdgsd

#define eps 1e-9
#define M_PI 3.141592653589793
#define bs 1000000009
#define bsize 350

using namespace std;

const int INF = 1e9;
const int N = 200551;

int n, k;
vector<int> g[N];
vector<int> gr[N];

int ans;
int used[N];

int go[N];

vector<int> get_cycle() {
  vector<int> res;
  int cur = 1;
  res.push_back(cur);
  cur = go[cur];
  while (cur != 1) {
    res.push_back(cur);
    cur = go[cur];
  }
  return res;
}

int indeg[N];
int on_cycle[N];
int D[N];

void trace_dep(int v, int dep) {
  D[v] = dep;
  used[v] = 1;
  for (int i = 0; i < gr[v].size(); i++) {
    int to = gr[v][i];
    trace_dep(to, dep + 1);
  }
}

int Res;

set<pair<int, int>> alive;
set<pair<int, int>>::iterator it;

int rem[N];

int GO(int v, int lev) {
  if (lev == 0)
    return v;
  return GO(go[v], lev - 1);
}

void update(int v, int steps) {

  if (alive.find(make_pair(D[v], v)) != alive.end()) {
    alive.erase(make_pair(D[v], v));
  }
  rem[v] = 1;
  for (int i = 0; i < gr[v].size(); i++) {
    int to = gr[v][i];
    if (rem[to])
      continue;
    update(to, steps - 1);
  }
}
int main() {
  // freopen("fabro.in","r",stdin);
  // freopen("fabro.out","w",stdout);
  // freopen("F:/in.txt", "r", stdin);
  // freopen("F:/output.txt", "w", stdout);
  ios_base::sync_with_stdio(0);
  // cin.tie(0);

  cin >> n >> k;
  for (int i = 1; i <= n; i++) {
    int to;
    cin >> go[i];
    if (i == 1)
      continue;

    indeg[go[i]]++;
    gr[go[i]].push_back(i);
  }

  if (go[1] != 1)
    ans++;
  go[1] = 1;

  trace_dep(1, 0);

  for (int i = 1; i <= n; i++) {
    alive.insert(make_pair(D[i], i));
  }

  while (true) {
    if (alive.size() == 0)
      break;
    it = alive.end();
    --it;
    int v = (*it).second;
    int WHR = GO(v, k - 1);
    if (go[WHR] != 1 && WHR != 1)
      ++ans;

    if (WHR == 1)
      update(WHR, 0);
    else
      update(WHR, 1);
  }

  cout << ans << endl;

  cin.get();
  cin.get();
  return 0;
}

Fixed submission:
/*
 */

#pragma comment(linker, "/STACK:16777216")
#define _CRT_SECURE_NO_WARNINGS

#include <algorithm>
#include <assert.h>
#include <complex>
#include <ctime>
#include <fstream>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <memory.h>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <vector>

#define y0 sdkfaslhagaklsldk
#define y1 aasdfasdfasdf
#define yn askfhwqriuperikldjk
#define j1 assdgsdgasghsf
#define tm sdfjahlfasfh
#define lr asgasgash
#define norm asdfasdgasdgsd

#define eps 1e-9
#define M_PI 3.141592653589793
#define bs 1000000009
#define bsize 350

using namespace std;

const int INF = 1e9;
const int N = 200551;

int n, k;
vector<int> g[N];
vector<int> gr[N];

int ans;
int used[N];

int go[N];

vector<int> get_cycle() {
  vector<int> res;
  int cur = 1;
  res.push_back(cur);
  cur = go[cur];
  while (cur != 1) {
    res.push_back(cur);
    cur = go[cur];
  }
  return res;
}

int indeg[N];
int on_cycle[N];
int D[N];

void trace_dep(int v, int dep) {
  D[v] = dep;
  used[v] = 1;
  for (int i = 0; i < gr[v].size(); i++) {
    int to = gr[v][i];
    trace_dep(to, dep + 1);
  }
}

int Res;

set<pair<int, int>> alive;
set<pair<int, int>>::iterator it;

int rem[N];

int GO(int v, int lev) {
  if (lev == 0)
    return v;
  if (v == 1)
    return v;
  return GO(go[v], lev - 1);
}

void update(int v, int steps) {

  if (alive.find(make_pair(D[v], v)) != alive.end()) {
    alive.erase(make_pair(D[v], v));
  }
  rem[v] = 1;
  for (int i = 0; i < gr[v].size(); i++) {
    int to = gr[v][i];
    if (rem[to])
      continue;
    update(to, steps - 1);
  }
}
int main() {
  // freopen("fabro.in","r",stdin);
  // freopen("fabro.out","w",stdout);
  // freopen("F:/in.txt", "r", stdin);
  // freopen("F:/output.txt", "w", stdout);
  ios_base::sync_with_stdio(0);
  // cin.tie(0);

  cin >> n >> k;
  for (int i = 1; i <= n; i++) {
    int to;
    cin >> go[i];
    if (i == 1)
      continue;

    indeg[go[i]]++;
    gr[go[i]].push_back(i);
  }

  if (go[1] != 1)
    ans++;
  go[1] = 1;

  trace_dep(1, 0);

  for (int i = 1; i <= n; i++) {
    alive.insert(make_pair(D[i], i));
  }

  while (true) {
    if (alive.size() == 0)
      break;
    it = alive.end();
    --it;
    int v = (*it).second;
    int WHR = GO(v, k - 1);
    if (go[WHR] != 1 && WHR != 1)
      ++ans;

    if (WHR == 1)
      update(WHR, 0);
    else
      update(WHR, 1);
  }

  cout << ans << endl;

  cin.get();
  cin.get();
  return 0;
}