Problem p02684 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
#include <cmath>
#include <iomanip>
#include <iterator>
#include <stdlib.h> /* srand, rand */
#include <time.h>   /* time */
using namespace std;

#define f first
#define s second

#define pb push_back
#define pp pop_back
#define mp make_pair

#define sz(x) (int)x.size()
#define sqr(x) ((x)*1ll * (x))
#define all(x) x.begin(), x.end()
#define clr(a, v) memset(a, v, sizeof(a))

#define rep(i, l, r) for (int i = (l); i <= (r); i++)
#define per(i, l, r) for (int i = (l); i >= (r); i--)
#define reflex ios_base::sync_with_stdio(false)

#define nl '\n'
#define ioi exit(0);

#define PI 3.14159265

typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;

typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
const int N = (int)2e5 + 7;
const int inf = (int)1e9 + 7;
const int mod = (int)1e9 + 7;
// const int mod = 1000000007;
const ll linf = (ll)1e18 + 7;

const int dx[] = {-1, 0, 1, 0, 1, -1, -1, 1};
const int dy[] = {0, 1, 0, -1, 1, -1, 1, -1};
const int kx[] = {-1, 1, 2, 2, 1, -1, -2, -2};
const int ky[] = {2, 2, 1, -1, -2, -2, -1, 1};

struct node {
  int x, y, z;
  mutable int w;
  node() { x = -1, y = -1, z = -1; }
  bool operator<(const node &nod) const {
    if (this->x == nod.x) {
      if (this->y == nod.y) {
        return this->z < nod.z;
      }
      return this->y < nod.y;
    }
    return this->x < nod.x;
  }
};

int a[N];
bool vis[N];
int main() {
  clr(vis, false);
  ll n, k;
  cin >> n >> k;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  int id = 1;
  while (!vis[id] && k > 0) {
    vis[id] = true;
    k--;
    id = a[id];
    //~ cout<<id<<"->";
  }
  if (k >= 0) {
    ll cnt = 1;
    int init = id;
    id = a[id];
    while (init != id) {
      id = a[id];
      cnt++;
    }

    k = k % cnt;
    //~ cout<<"k:"<< k <<endl;
    while (k) {
      k--;
      id = a[id];
    }
  }
  cout << id << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
#include <cmath>
#include <iomanip>
#include <iterator>
#include <stdlib.h> /* srand, rand */
#include <time.h>   /* time */
using namespace std;

#define f first
#define s second

#define pb push_back
#define pp pop_back
#define mp make_pair

#define sz(x) (int)x.size()
#define sqr(x) ((x)*1ll * (x))
#define all(x) x.begin(), x.end()
#define clr(a, v) memset(a, v, sizeof(a))

#define rep(i, l, r) for (int i = (l); i <= (r); i++)
#define per(i, l, r) for (int i = (l); i >= (r); i--)
#define reflex ios_base::sync_with_stdio(false)

#define nl '\n'
#define ioi exit(0);

#define PI 3.14159265

typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;

typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
const int N = (int)2e5 + 7;
const int inf = (int)1e9 + 7;
const int mod = (int)1e9 + 7;
// const int mod = 1000000007;
const ll linf = (ll)1e18 + 7;

const int dx[] = {-1, 0, 1, 0, 1, -1, -1, 1};
const int dy[] = {0, 1, 0, -1, 1, -1, 1, -1};
const int kx[] = {-1, 1, 2, 2, 1, -1, -2, -2};
const int ky[] = {2, 2, 1, -1, -2, -2, -1, 1};

struct node {
  int x, y, z;
  mutable int w;
  node() { x = -1, y = -1, z = -1; }
  bool operator<(const node &nod) const {
    if (this->x == nod.x) {
      if (this->y == nod.y) {
        return this->z < nod.z;
      }
      return this->y < nod.y;
    }
    return this->x < nod.x;
  }
};

int a[N];
bool vis[N];
int main() {
  clr(vis, false);
  ll n, k;
  cin >> n >> k;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  int id = 1;
  while (!vis[id] && k > 0) {
    vis[id] = true;
    k--;
    id = a[id];
    //~ cout<<id<<"->";
  }
  if (k > 0) {
    ll cnt = 1;
    int init = id;
    id = a[id];
    while (init != id) {
      id = a[id];
      cnt++;
    }

    k = k % cnt;
    //~ cout<<"k:"<< k <<endl;
    while (k) {
      k--;
      id = a[id];
    }
  }
  cout << id << endl;
  return 0;
}