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

#include <algorithm>
#include <bits/stdc++.h>
#include <climits>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <iostream>
#include <math.h>
using namespace __gnu_pbds;
using namespace std;
#define ll long long
#define str string
#define deb_arr(arr, n)                                                        \\
  for (ll i = 0; i < (n); i++) {                                               \\
    cout << (arr)[i] << " ";                                                   \\
  }
#define in_arr0(arr, n)                                                        \\
  for (ll i = 0; i < (n); i++) {                                               \\
    cin >> (arr)[i];                                                           \\
  }
#define in_arr1(arr, n)                                                        \\
  for (ll i = 1; i <= (n); i++) {                                              \\
    cin >> (arr)[i];                                                           \\
  }
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;
#define vll vector<ll>
#define fr(n) for (i = 0; i < (n); i++)
#define fr1(n) for (i = 1; i <= (n); i++)
#define pll pair<ll, ll>
#define f first
#define s second
#define vpll vector<pll>
#define ump unordered_map
#define YES cout << "YES" << endl
#define NO cout << "NO" << endl
#define IN(a, b) ((a).find(b) != (a).end())
#define pb push_back
#define mkp make_pair
#define PI M_PI
#define INF 1e18
#define deb(x)                                                                 \\
  if (ON)                                                                      \\
  cout << #x << " : " << x << endl
#define fast ios_base::sync_with_stdio(false)
#define io cin.tie(NULL), cout.tie(NULL)
#define ON !false
ll dx[] = {0, -1, 0, 1};
ll dy[] = {1, 0, -1, 0};
const ll mod = 1e9 + 7;
inline ll add(ll a, ll b) { return ((a % mod) + (b % mod) + mod) % mod; }
inline ll mul(ll a, ll b) { return ((a % mod) * (b % mod) + mod) % mod; }
void pause(ll id) {
  cout << "Waiting " << id << endl;
  cin >> id;
}

int mxn = 2e5 + 10;

void init() { return; }

signed main() {
  fast;
  io;
  if (!true) {
    freopen("in.txt", "r", stdin);
    freopen("out.txt", "w", stdout);
  }
  init();
  ll i, j;
  ll n, k;
  cin >> n >> k;
  ll arr[n + 1];
  in_arr1(arr, n);
  ll cnt = 0;
  vll vis(n + 1, 0);
  vll ans;
  i = 1;
  map<ll, ll> mp;
  while (true) {
    if (vis[i]) {
      break;
    }
    mp[i] = cnt;
    cnt++;
    vis[i] = 1;
    ans.pb(i);
    i = arr[i];
  }
  ll ok;
  k++;
  i = mp[i];

  if (k <= i) {
    cout << ans[k - 1] << endl;
    return 0;
  }
  k -= i;
  ll sz = (ans.size() - i);
  k = k % sz;
  cout << ans[i + (k - 1 + sz) % sz] << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bits/stdc++.h>
#include <climits>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <iostream>
#include <math.h>
using namespace __gnu_pbds;
using namespace std;
#define ll long long
#define str string
#define deb_arr(arr, n)                                                        \\
  for (ll i = 0; i < (n); i++) {                                               \\
    cout << (arr)[i] << " ";                                                   \\
  }
#define in_arr0(arr, n)                                                        \\
  for (ll i = 0; i < (n); i++) {                                               \\
    cin >> (arr)[i];                                                           \\
  }
#define in_arr1(arr, n)                                                        \\
  for (ll i = 1; i <= (n); i++) {                                              \\
    cin >> (arr)[i];                                                           \\
  }
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;
#define vll vector<ll>
#define fr(n) for (i = 0; i < (n); i++)
#define fr1(n) for (i = 1; i <= (n); i++)
#define pll pair<ll, ll>
#define f first
#define s second
#define vpll vector<pll>
#define ump unordered_map
#define YES cout << "YES" << endl
#define NO cout << "NO" << endl
#define IN(a, b) ((a).find(b) != (a).end())
#define pb push_back
#define mkp make_pair
#define PI M_PI
#define INF 1e18
#define deb(x)                                                                 \\
  if (ON)                                                                      \\
  cout << #x << " : " << x << endl
#define fast ios_base::sync_with_stdio(false)
#define io cin.tie(NULL), cout.tie(NULL)
#define ON !false
ll dx[] = {0, -1, 0, 1};
ll dy[] = {1, 0, -1, 0};
const ll mod = 1e9 + 7;
inline ll add(ll a, ll b) { return ((a % mod) + (b % mod) + mod) % mod; }
inline ll mul(ll a, ll b) { return ((a % mod) * (b % mod) + mod) % mod; }
void pause(ll id) {
  cout << "Waiting " << id << endl;
  cin >> id;
}

int mxn = 2e5 + 10;

void init() { return; }

signed main() {
  fast;
  io;
  if (!true) {
    freopen("in.txt", "r", stdin);
    freopen("out.txt", "w", stdout);
  }
  init();
  ll i, j;
  ll n, k;
  cin >> n >> k;
  ll arr[n + 1];
  in_arr1(arr, n);
  ll cnt = 0;
  vll vis(n + 1, 0);
  vll ans;
  i = 1;
  while (true) {
    if (vis[i]) {
      break;
    }
    cnt++;
    vis[i] = 1;
    ans.pb(i);
    i = arr[i];
  }
  ll ok;
  k++;
  while (mxn-- && ans.size() > 0 && ans.front() != i) {
    ok = ans.front();
    ans.erase(ans.begin());
    k--;
    if (k == 0) {
      break;
    }
  }
  if (k == 0) {
    cout << ok << endl;
    return 0;
  }

  ll sz = ans.size();
  k = k % sz;
  cout << ans[(k - 1 + sz) % sz] << endl;

  return 0;
}