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

#include <bits/stdc++.h>
#define rep(i, n) for (ll i = 0; i < n; i++)
#define repr(i, n) for (ll i = n; i >= 0; i--)
#define reprr(i, m, n) for (ll i = m; i < n; i++)
#define inf 2e9
#define all(v) v.begin(), v.end()
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
typedef vector<vll> vvll;

struct UnionFind {
  vll par;

  UnionFind(ll N) : par(N) { rep(i, N) par[i] = i; }

  ll root(ll x) {
    if (par[x] == x)
      return x;
    return par[x] = root(par[x]);
  }

  void unite(ll x, ll y) {
    ll rx = root(x);
    ll ry = root(y);
    if (rx == ry)
      return;
    par[rx] = ry;
  }

  bool same(ll x, ll y) {
    ll rx = root(x);
    ll ry = root(y);
    return rx == ry;
  }
};

int main() {
  ll N, K;
  cin >> N >> K;
  vll A(N);
  rep(i, N) cin >> A[i], A[i]--;
  vll visited(N, -1);
  vll path{0};
  visited[0] = 0;
  ll index = 0;
  rep(i, N) {
    // i: how many times takahashi teleport
    if (i == K) {
      cout << index + 1 << endl;
      return 0;
    }
    if (index == A[index]) {
      cout << index + 1 << endl;
      return 0;
    }

    // cout << "i: " << i << endl;

    // do teleport
    index = A[index];

    if (visited[index] >= 0) {
      // find loop
      ll base = visited[index];
      ll loop_range = i + 1 - base;
      // cout << "i: " << i << endl;
      // cout << "base: " << base << endl;
      // cout << "loop_range: " << loop_range << endl;
      cout << path[((K - base) % (loop_range)) + base] + 1 << endl;
      return 0;
    } else {
      visited[index] = i + 1;
      path.push_back(index);
      // cout << "visited:: ";
      // for (auto elem : visited) cout << elem << " "; cout << endl;
      // cout << "path:: ";
      // for (auto elem : path) cout << elem << " "; cout << endl;
    }
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (ll i = 0; i < n; i++)
#define repr(i, n) for (ll i = n; i >= 0; i--)
#define reprr(i, m, n) for (ll i = m; i < n; i++)
#define inf 2e9
#define all(v) v.begin(), v.end()
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
typedef vector<vll> vvll;

struct UnionFind {
  vll par;

  UnionFind(ll N) : par(N) { rep(i, N) par[i] = i; }

  ll root(ll x) {
    if (par[x] == x)
      return x;
    return par[x] = root(par[x]);
  }

  void unite(ll x, ll y) {
    ll rx = root(x);
    ll ry = root(y);
    if (rx == ry)
      return;
    par[rx] = ry;
  }

  bool same(ll x, ll y) {
    ll rx = root(x);
    ll ry = root(y);
    return rx == ry;
  }
};

int main() {
  ll N, K;
  cin >> N >> K;
  vll A(N);
  rep(i, N) cin >> A[i], A[i]--;
  ll cnt = 0;
  vll visited{0};
  ll index = 0;
  rep(i, N) {
    if (cnt == K) {
      cout << index + 1 << endl;
      return 0;
    }
    if (index == A[index]) {
      cout << index + 1 << endl;
      return 0;
    }

    // cout << "i: " << i << endl;

    index = A[index];
    auto iter = find(all(visited), index);

    if (iter != visited.end()) {
      // find loop
      ll index = iter - visited.begin();
      ll loop_range = visited.size() - index;
      // cout << "index: " << index << endl;
      // cout << "loop_range: " << loop_range << endl;
      cout << visited[((K - index) % (loop_range)) + index] + 1 << endl;
      return 0;
    } else {
      visited.push_back(index);
      // cout << "visited:: ";
      // for (auto elem : visited) cout << elem << " "; cout << endl;
    }
    cnt++;
  }

  return 0;
}