Problem p02684 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;

#define pii pair<int, int>
#define vi vector<int>
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define x first
#define y second
#define sz(x) (x).size()
#define all(x) (x).begin(), (x).end()  // Forward traversal
#define rall(x) (x).rbegin, (x).rend() // reverse traversal

#define ll long long
#define ld long double

const ld EPS = 10e-9;
const int INF = 1e9 + 5;
const long long INFLL = (ll)INF * (ll)INF;
const ll MOD = 1e9 + 7;
const ld PI = 3.14159265358979323846;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  for (int i = 0; i < (int)v.size(); ++i) {
    os << v[i];
    if (i != (int)v.size() - 1)
      os << ", ";
  }
  os << "]\n";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\n";
  return os;
}

template <typename T>
ostream &operator<<(ostream &os, const unordered_set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const unordered_map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const pair<T, S> &v) {
  os << "(";
  os << v.first << ", " << v.second << ")";
  os << "\n";
  return os;
}

#define DEBUG

#ifdef DEBUG
#define debug(...) __f(#__VA_ARGS__, __VA_ARGS__)
template <typename Arg1> void __f(const char *name, Arg1 &&arg1) {
  cerr << name << " : " << arg1;
}
template <typename Arg1, typename... Args>
void __f(const char *names, Arg1 &&arg1, Args &&...args) {
  const char *comma = strchr(names + 1, ',');
  cerr.write(names, comma - names) << " : " << arg1 << " | ";
  __f(comma + 1, args...);
}
#else
#define debug(...)
#endif

// __gcd(x, y)

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);

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

  vector<ll> a(N);
  for (int i = 0; i < N; i++)
    cin >> a[i];

  ll l = a[0] - 1;
  ll r = a[a[0] - 1] - 1;
  while (l != r) {
    l = a[l] - 1;
    r = a[a[r] - 1] - 1;
  }
  l = 0;
  while (l != r) {
    l = a[l] - 1;
    r = a[r] - 1;
  }
  // floyd warshall cycle detection algorithm ends here to find start of cycle

  ll cycleStar = l; // starting point of cycle
  ll cnt = 0;       // number of elements before the starting point of cycle
  l = 0;
  while (l != cycleStar) {
    cnt++;
    l = a[l] - 1;
  }

  ll pos = cnt; // number of elements before the starting point of cycle
  ll c = 1;     // length of cycle(num of elements in it)
  l = cycleStar;
  while (a[l] - 1 != cycleStar) {
    l = a[l] - 1;
    c++;
  }

  ll i = 0;
  unordered_set<ll> visited;
  visited.insert(0);
  ll pos2 = -1;
  ll c2 = 1;
  int lenBefore = 0;
  while (true) {
    if (visited.count(a[i] - 1)) {
      i = a[i] - 1;
      pos2 = 0;
      while (pos2 != i) {
        pos2 = a[pos2] - 1;
        lenBefore++;
      }
      pos2 = i;
      while (pos2 != a[i] - 1) {
        c2++;
        i = a[i] - 1;
      }
      break;
    }
    visited.insert(a[i] - 1);
    i = a[i] - 1;
  }

  if (pos != lenBefore) {
    throw runtime_error(to_string(pos) + " my: " + to_string(lenBefore));
  }
  //    if(pos != pos2)
  //    {
  //        throw runtime_error(to_string(pos) + " my: " + to_string(pos2));
  //    }

  ll ans = 0;
  while (lenBefore != 0 && K != 0) {
    ans = a[ans] - 1;
    pos--;
    K--;
  }
  if (K == 0) {
    cout << ans + 1 << endl;
    return 0;
  }

  ll n = K % c;
  while (n) {
    ans = a[ans] - 1;
    n--;
  }

  cout << ans + 1 << endl;
  //    if(K <= pos)
  //    {
  //        ll ans = 0;
  //        while(K)
  //        {
  //            K--;
  //            ans = a[ans]-1;
  //        }
  //        cout << ans+1 << endl;
  //    }
  //    else {
  //        ll ans = 0;
  //        int pos_c = pos;
  //        while(pos_c)
  //        {
  //            pos_c--;
  //            ans = a[ans]-1;
  //        }
  //        int n = (K - pos) % c;
  //        while(n)
  //        {
  //            n--;
  //            ans = a[ans]-1;
  //        }
  //        cout << ans+1 << endl;
  //    }
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

#define pii pair<int, int>
#define vi vector<int>
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define x first
#define y second
#define sz(x) (x).size()
#define all(x) (x).begin(), (x).end()  // Forward traversal
#define rall(x) (x).rbegin, (x).rend() // reverse traversal

#define ll long long
#define ld long double

const ld EPS = 10e-9;
const int INF = 1e9 + 5;
const long long INFLL = (ll)INF * (ll)INF;
const ll MOD = 1e9 + 7;
const ld PI = 3.14159265358979323846;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  for (int i = 0; i < (int)v.size(); ++i) {
    os << v[i];
    if (i != (int)v.size() - 1)
      os << ", ";
  }
  os << "]\n";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\n";
  return os;
}

template <typename T>
ostream &operator<<(ostream &os, const unordered_set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const unordered_map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const pair<T, S> &v) {
  os << "(";
  os << v.first << ", " << v.second << ")";
  os << "\n";
  return os;
}

#define DEBUG

#ifdef DEBUG
#define debug(...) __f(#__VA_ARGS__, __VA_ARGS__)
template <typename Arg1> void __f(const char *name, Arg1 &&arg1) {
  cerr << name << " : " << arg1;
}
template <typename Arg1, typename... Args>
void __f(const char *names, Arg1 &&arg1, Args &&...args) {
  const char *comma = strchr(names + 1, ',');
  cerr.write(names, comma - names) << " : " << arg1 << " | ";
  __f(comma + 1, args...);
}
#else
#define debug(...)
#endif

// __gcd(x, y)

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);

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

  vector<ll> a(N);
  for (int i = 0; i < N; i++)
    cin >> a[i];

  ll l = a[0] - 1;
  ll r = a[a[0] - 1] - 1;
  while (l != r) {
    l = a[l] - 1;
    r = a[a[r] - 1] - 1;
  }
  l = 0;
  while (l != r) {
    l = a[l] - 1;
    r = a[r] - 1;
  }
  // floyd warshall cycle detection algorithm ends here to find start of cycle

  ll cycleStar = l; // starting point of cycle
  ll cnt = 0;       // number of elements before the starting point of cycle
  l = 0;
  while (l != cycleStar) {
    cnt++;
    l = a[l] - 1;
  }

  ll pos = cnt; // number of elements before the starting point of cycle
  ll c = 1;     // length of cycle(num of elements in it)
  l = cycleStar;
  while (a[l] - 1 != cycleStar) {
    l = a[l] - 1;
    c++;
  }

  ll i = 0;
  unordered_set<ll> visited;
  visited.insert(0);
  ll pos2 = -1;
  ll c2 = 1;
  int lenBefore = 0;
  while (true) {
    if (visited.count(a[i] - 1)) {
      i = a[i] - 1;
      pos2 = 0;
      while (pos2 != i) {
        pos2 = a[pos2] - 1;
        lenBefore++;
      }
      pos2 = i;
      while (pos2 != a[i] - 1) {
        c2++;
        i = a[i] - 1;
      }
      break;
    }
    visited.insert(a[i] - 1);
    i = a[i] - 1;
  }

  if (pos != lenBefore) {
    throw runtime_error(to_string(pos) + " my: " + to_string(lenBefore));
  }
  //    if(pos != pos2)
  //    {
  //        throw runtime_error(to_string(pos) + " my: " + to_string(pos2));
  //    }

  ll ans = 0;
  while (lenBefore != 0 && K != 0) {
    ans = a[ans] - 1;
    lenBefore--;
    K--;
  }
  if (K == 0) {
    cout << ans + 1 << endl;
    return 0;
  }

  ll n = K % c;
  while (n) {
    ans = a[ans] - 1;
    n--;
  }

  cout << ans + 1 << endl;
  //    if(K <= pos)
  //    {
  //        ll ans = 0;
  //        while(K)
  //        {
  //            K--;
  //            ans = a[ans]-1;
  //        }
  //        cout << ans+1 << endl;
  //    }
  //    else {
  //        ll ans = 0;
  //        int pos_c = pos;
  //        while(pos_c)
  //        {
  //            pos_c--;
  //            ans = a[ans]-1;
  //        }
  //        int n = (K - pos) % c;
  //        while(n)
  //        {
  //            n--;
  //            ans = a[ans]-1;
  //        }
  //        cout << ans+1 << endl;
  //    }
  return 0;
}