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

#include <bits/stdc++.h>
#define int int64_t
#define REP(i, n) for (int i = 0; (i) < (int)(n); ++(i))
#define REP3(i, m, n) for (int i = (m); (i) < (int)(n); ++(i))
#define REP_R(i, n) for (int i = (int)(n)-1; (i) >= 0; --(i))
#define REP3R(i, m, n) for (int i = (int)(n)-1; (i) >= (int)(m); --(i))
#define REPP(i, n) for (int i = 1; (i) < (int)(n + 1); ++(i))
#define REPR(i, n) for (int i = (int)(n)-1; (i) >= 0; --(i))
#define REPRP(i, n) for (int i = (int)(n); (i) >= 0; --(i))
#define FOR(i, m, n) for (int i = m; i < n; i++)
#define FORP(i, m, n) for (int i = m; i < n + 1; i++)
#define FORR(i, a, b) for (int i = (b)-1; i >= (a); i--)
#define ALL(v) v.begin(), v.end()
#define SZ(a) int((a).size())
#define SORT(c) sort((c).begin(), (c).end())
#define SORTR(c) sort((c).begin(), (c).end(), greater<int>());
#define BIT(n) (1LL << (n))
#define PCNT(x) __builtin_popcountll(x)
#define GCD(x, y) __gcd(x, y)
#define LCM(x, y) x / __gcd(x, y) * y
#define PB push_back
#define MP make_pair
#define endl "\\n"
#define pa cout << ans << endl
#define p(x) cout << x << endl
#define pall(x)                                                                \\
  for (auto iter : x)                                                          \\
  p(iter)
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define yn(x) cout << (x ? "Yes" : "No") << endl
using namespace std;
typedef long double LD;
typedef vector<int> VI;
typedef vector<vector<int>> VVI;
typedef vector<string> VS;
typedef pair<int, int> PII;
typedef vector<PII> VPII;
using Graph = vector<vector<int>>;
using WeightedGraph = vector<vector<pair<int, int>>>;
const int INF = 2000000000ll;
const long double PI = acosl(-1);
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

int power(int a, int b) {
  return b ? power(a * a, b / 2) * (b % 2 ? a : 1) : 1;
}
int perm(int N, int K) {
  int x = 1;
  for (int i = 1; i <= K; i++) {
    x = x * (N - i + 1);
  }
  return x;
}
int factorial(int N) { return perm(N, N); }
int comb(int N, int K) {
  return factorial(N) / factorial(K) / factorial(N - K);
}
int hcomb(int N, int K) { return comb(N + K - 1, N - 1); }

int dtoi(double d) { return round(d); }

int isprime(int n) {
  if (n < 2)
    return 0;
  for (int i = 2; i * i <= n; ++i) {
    if (n % i == 0)
      return 0;
  }
  return 1;
}

int mystoi(string S) {
  int num = 0;
  for (int i = (int)S.length() - 1; i >= 0; i--) {
    num += (S[i] - '0') * power(10, S.length() - 1 - i);
  }
  return num;
}

template <class T> int SUM(T a) {
  int sum = 0;
  for (auto iter : a)
    sum += iter;
  return sum;
}

void solve(int N, int K, const vector<int> &A);

signed main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(nullptr);
  cout << fixed << setprecision(10);
  int N, K;
  cin >> N;
  vector<int> A(N);
  cin >> K;
  REP(i, N) {
    cin >> A[i];
    A[i]--;
  }
  solve(N, K, A);
  return 0;
}

void solve(int N, int K, const vector<int> &A) {
  int ans = 0;
  int now = 0;
  int cnt = 1;
  int start = -1;
  unordered_map<int, int> zumi;
  while (1) {
    if (zumi[now] == 1) {
      start = now;
      break;
    }
    if (start != -1)
      break;
    zumi[now] = 1;
    now = A[now];
  }
  now = A[start];
  while (now != start) {
    cnt++;
    now = A[now];
  }

  now = 0;
  int temp = 0;
  while (now != start) {
    temp++;
    now = A[now];
  }

  now = 0;
  if (K <= temp) {
    REP(i, K) now = A[now];
    p(now + 1);
    return;
  } else {
    K -= temp;
    K %= cnt;
    now = start;
    REP(i, K) now = A[now];
    p(now + 1);
    return;
  }
}

/*
cd ~/Desktop/Programming/AtCoder/contest/abc167/abc167_d && sh test.sh
cd ~/Desktop/Programming/AtCoder/contest/abc167/abc167_d && sh etest.sh
cd ~/Desktop/Programming/AtCoder/contest/abc167/abc167_d && sh submit.sh
*/


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define int int64_t
#define REP(i, n) for (int i = 0; (i) < (int)(n); ++(i))
#define REP3(i, m, n) for (int i = (m); (i) < (int)(n); ++(i))
#define REP_R(i, n) for (int i = (int)(n)-1; (i) >= 0; --(i))
#define REP3R(i, m, n) for (int i = (int)(n)-1; (i) >= (int)(m); --(i))
#define REPP(i, n) for (int i = 1; (i) < (int)(n + 1); ++(i))
#define REPR(i, n) for (int i = (int)(n)-1; (i) >= 0; --(i))
#define REPRP(i, n) for (int i = (int)(n); (i) >= 0; --(i))
#define FOR(i, m, n) for (int i = m; i < n; i++)
#define FORP(i, m, n) for (int i = m; i < n + 1; i++)
#define FORR(i, a, b) for (int i = (b)-1; i >= (a); i--)
#define ALL(v) v.begin(), v.end()
#define SZ(a) int((a).size())
#define SORT(c) sort((c).begin(), (c).end())
#define SORTR(c) sort((c).begin(), (c).end(), greater<int>());
#define BIT(n) (1LL << (n))
#define PCNT(x) __builtin_popcountll(x)
#define GCD(x, y) __gcd(x, y)
#define LCM(x, y) x / __gcd(x, y) * y
#define PB push_back
#define MP make_pair
#define endl "\\n"
#define pa cout << ans << endl
#define p(x) cout << x << endl
#define pall(x)                                                                \\
  for (auto iter : x)                                                          \\
  p(iter)
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define yn(x) cout << (x ? "Yes" : "No") << endl
using namespace std;
typedef long double LD;
typedef vector<int> VI;
typedef vector<vector<int>> VVI;
typedef vector<string> VS;
typedef pair<int, int> PII;
typedef vector<PII> VPII;
using Graph = vector<vector<int>>;
using WeightedGraph = vector<vector<pair<int, int>>>;
const int INF = 2000000000ll;
const long double PI = acosl(-1);
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

int power(int a, int b) {
  return b ? power(a * a, b / 2) * (b % 2 ? a : 1) : 1;
}
int perm(int N, int K) {
  int x = 1;
  for (int i = 1; i <= K; i++) {
    x = x * (N - i + 1);
  }
  return x;
}
int factorial(int N) { return perm(N, N); }
int comb(int N, int K) {
  return factorial(N) / factorial(K) / factorial(N - K);
}
int hcomb(int N, int K) { return comb(N + K - 1, N - 1); }

int dtoi(double d) { return round(d); }

int isprime(int n) {
  if (n < 2)
    return 0;
  for (int i = 2; i * i <= n; ++i) {
    if (n % i == 0)
      return 0;
  }
  return 1;
}

int mystoi(string S) {
  int num = 0;
  for (int i = (int)S.length() - 1; i >= 0; i--) {
    num += (S[i] - '0') * power(10, S.length() - 1 - i);
  }
  return num;
}

template <class T> int SUM(T a) {
  int sum = 0;
  for (auto iter : a)
    sum += iter;
  return sum;
}

void solve(int N, int K, const vector<int> &A);

signed main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(nullptr);
  cout << fixed << setprecision(10);
  int N, K;
  cin >> N;
  vector<int> A(N);
  cin >> K;
  REP(i, N) {
    cin >> A[i];
    A[i]--;
  }
  solve(N, K, A);
  return 0;
}

void solve(int N, int K, const vector<int> &A) {
  int ans = 0;
  int now = 0;
  int cnt = 1;
  int start = -1;
  VI zumi;
  while (1) {
    for (int i : zumi)
      if (i == now) {
        start = i;
        break;
      }
    if (start != -1)
      break;
    zumi.PB(now);
    now = A[now];
  }
  now = A[start];
  while (now != start) {
    cnt++;
    now = A[now];
  }

  now = 0;
  int temp = 0;
  while (now != start) {
    temp++;
    now = A[now];
  }

  now = 0;
  if (K <= temp) {
    REP(i, K) now = A[now];
    p(now + 1);
    return;
  } else {
    K -= temp;
    K %= cnt;
    now = start;
    REP(i, K) now = A[now];
    p(now + 1);
    return;
  }
}

/*
cd ~/Desktop/Programming/AtCoder/contest/abc167/abc167_d && sh test.sh
cd ~/Desktop/Programming/AtCoder/contest/abc167/abc167_d && sh etest.sh
cd ~/Desktop/Programming/AtCoder/contest/abc167/abc167_d && sh submit.sh
*/