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

#ifdef _MSC_VER
#include "stdc++.h"
#endif
#ifdef __GNUC__
#include <bits/stdc++.h>
#endif
using namespace std;

using ll = long long;

// vector
using vi = vector<int>;
using vii = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;

// constant
const double EPS = 1e-10;

#define REP(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, s, e) for (ll i = s; i < e; i++)
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()

// debug
#define DUMP(x) cerr << #x << " = " << (x) << endl;
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

// vector
// template<typename T> vector<vector<T>> vv(ll H, ll W, T v) { return
// vector<vector<T>> vec(H, vector<T>(W, v)); }

// view vector
template <typename T> void view(T e) { cout << e << endl; }
template <typename T> void view(const vector<T> &v) {
  for (const auto &e : v) {
    cout << e << " ";
  }
  cout << endl;
}
template <typename T> void view(const vector<vector<T>> &vv) {
  for (const auto &v : vv) {
    view(v);
  }
}

// DP
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;
}

// ***************************************

int main() {

  ll N, K;
  cin >> N >> K;

  vll A(N);
  REP(i, N) {
    cin >> A[i];
    A[i]--;
  }

  vvll visited(N);
  vll rireki = {0};

  ll cur = 0;
  bool flag = false;
  REP(i, K) {
    cur = A[cur];
    if (!visited[cur].empty()) {
      flag = true;
      break;
    }
    rireki.push_back(cur);
    visited[cur].push_back(i + 1);
  }

  if (!flag) {
    cout << cur + 1 << endl;
    return 0;
  }

  ll second = rireki.size();
  ll first = visited[cur][0];

  ll idx = (K - first) % (second - first);

  cout << rireki[first + idx] + 1 << endl;

  return 0;
}

Error: TLE

Fixed submission:

#ifdef _MSC_VER
#include "stdc++.h"
#endif
#ifdef __GNUC__
#include <bits/stdc++.h>
#endif
using namespace std;

using ll = long long;

// vector
using vi = vector<int>;
using vii = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;

// constant
const double EPS = 1e-10;

#define REP(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, s, e) for (ll i = s; i < e; i++)
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()

// debug
#define DUMP(x) cerr << #x << " = " << (x) << endl;
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

// vector
// template<typename T> vector<vector<T>> vv(ll H, ll W, T v) { return
// vector<vector<T>> vec(H, vector<T>(W, v)); }

// view vector
template <typename T> void view(T e) { cout << e << endl; }
template <typename T> void view(const vector<T> &v) {
  for (const auto &e : v) {
    cout << e << " ";
  }
  cout << endl;
}
template <typename T> void view(const vector<vector<T>> &vv) {
  for (const auto &v : vv) {
    view(v);
  }
}

// DP
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;
}

// ***************************************

int main() {

  ll N, K;
  cin >> N >> K;

  vll A(N);
  REP(i, N) {
    cin >> A[i];
    A[i]--;
  }

  map<ll, ll> rireki;
  vector<ll> rireki2 = {0};
  rireki[0] = 0;
  set<ll> visited = {0};

  ll cur = 0;
  bool flag = false;
  REP(i, K) {
    cur = A[cur];
    if (binary_search(ALL(visited), cur)) {
      flag = true;
      break;
    }
    rireki[cur] = i + 1;
    rireki2.push_back(cur);
    visited.insert(cur);
  }

  if (!flag) {
    cout << cur + 1 << endl;
    return 0;
  }

  ll second = rireki2.size();
  ll first = rireki[cur];

  ll idx = (K - first) % (second - first);

  cout << rireki2[first + idx] + 1 << endl;

  return 0;
}