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

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

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP0(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define RREP0(i, n) for (int i = (n)-1; i >= 0; --i)
#define REP1(i, n) for (int i = 1, i##_len = (n); i <= i##_len; ++i)
#define RREP1(i, n) for (int i = (n); i >= 1; --i)
#define pow10(n) int(1e##n + n)
typedef long long LL;
typedef pair<int, int> P;

const int INTINF = int(1e9) + 1;
const LL LLINF = LL(1e18) + 1;
long double eps = 1.0E-14;

int dist[2 * pow10(5)];
int nextOf[2 * pow10(5)];

void solve() {
  int N;
  LL K;
  cin >> N >> K;

  REP0(i, N) {
    dist[i] = INTINF;
    cin >> nextOf[i];
    nextOf[i]--;
  }

  if (nextOf[0] == 0) {
    std::cout << 1 << endl;
    return;
  }

  int cur = 0;
  dist[0] = 0;
  LL cost_till_loop = 0, l_size = 0, loop_start = 0, nxt = 0;

  int count = 0;

  while (true) {
    nxt = nextOf[cur];
    // std::cout << "->" << nxt << endl;
    count++;
    if (K == count) {
      std::cout << nxt + 1 << endl;
      return;
    }
    if (nxt == cur) {
      std::cout << nxt + 1 << endl;
      return;
    }
    if (dist[nxt] < dist[cur] + 1) {
      loop_start = nxt;
      cost_till_loop = dist[nxt];
      l_size = dist[cur] - dist[nxt] + 1;
      break;
    } else {
      dist[nxt] = dist[cur] + 1;
      cur = nxt;
    }
  }

  K -= cost_till_loop;
  LL inner = K % l_size;
  // if (inner == 0)
  // {
  //     std::cout << loop_start + 1 << endl;
  //     return;
  // }

  cur = loop_start;
  int cnt = 0;
  while (true) {
    cur = nextOf[cur];
    cnt++;
    if (cnt % l_size == inner) {
      std::cout << cur + 1 << endl;
      return;
    }
  }
}

int main(int argc, char const *argv[]) {
  cin.tie(0);
  ios::sync_with_stdio(false);
  // std::cout << std::fixed << std::setprecision(20);
  solve();
  return 0;
}


Error: TLE

Fixed submission:

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

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP0(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define RREP0(i, n) for (int i = (n)-1; i >= 0; --i)
#define REP1(i, n) for (int i = 1, i##_len = (n); i <= i##_len; ++i)
#define RREP1(i, n) for (int i = (n); i >= 1; --i)
#define pow10(n) int(1e##n + n)
typedef long long LL;
typedef pair<int, int> P;

const int INTINF = int(1e9) + 1;
const LL LLINF = LL(1e18) + 1;
long double eps = 1.0E-14;

int dist[2 * pow10(5)];
int nextOf[2 * pow10(5)];

void solve() {
  int N;
  LL K;
  cin >> N >> K;

  REP0(i, N) {
    dist[i] = INTINF;
    cin >> nextOf[i];
    nextOf[i]--;
  }

  int cur = 0;
  dist[0] = 0;
  LL cost_till_loop = 0, l_size = 0, loop_start = 0, nxt = 0;

  int count = 0;

  while (true) {
    nxt = nextOf[cur];
    // std::cout << "->" << nxt << endl;
    count++;
    if (K == count) {
      std::cout << nxt + 1 << endl;
      return;
    }
    if (dist[nxt] < dist[cur]) {
      loop_start = nxt;
      cost_till_loop = dist[nxt];
      l_size = dist[cur] - dist[nxt] + 1;
      break;
    } else {
      dist[nxt] = dist[cur] + 1;
      cur = nxt;
    }
  }

  K -= cost_till_loop;
  LL inner = K % l_size;
  // if (inner == 0)
  // {
  //     std::cout << loop_start + 1 << endl;
  //     return;
  // }

  cur = loop_start;
  int cnt = 0;
  while (true) {
    cur = nextOf[cur];
    cnt++;
    if (cnt % l_size == inner) {
      std::cout << cur + 1 << endl;
      return;
    }
  }
}

int main(int argc, char const *argv[]) {
  cin.tie(0);
  ios::sync_with_stdio(false);
  // std::cout << std::fixed << std::setprecision(20);
  solve();
  return 0;
}