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

#include <bits/stdc++.h>
using namespace std;
using ll = long long int;
using ld = long double;
using VI = vector<ll>;
using VD = vector<ld>;
using VVI = vector<VI>;
using VC = vector<char>;
using VVC = vector<VC>;
using VS = vector<string>;
using PLL = pair<ll, ll>;
using PLD = pair<ld, ld>;
using VPLL = vector<PLL>;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
#define REPD(i, n) for (ll i = (ll)(n)-1; i >= 0; i--)
#define ALL(x) (x).begin(), (x).end()
#define ALLR(x) (x).rbegin(), (x).rend()
#define SZ(x) ((ll)(x).size())
#define MAX(x) *max_element((x).begin(), (x).end())
#define MIN(x) *min_element((x).begin(), (x).end())
#define SORTR(x) sort((x).rbegin(), (x).rend())
#define SORT(x) sort((x).begin(), (x).end())
#define SUM(x) accumulate((x).begin(), (x).end(), 0)
#define FILL(x, a) fill(x.begin(), x.end(), a)
#define EACH(i, x)                                                             \\
  for (typeof((x).begin()) i = (x).begin(); i != (x).end(); ++i)
#define print(x) cout << x << "\\n"
#define EXIST(v, x) (std::find(v.begin(), v.end(), x) != v.end())

const ll INF = 1e18;
const ld EPS = 1e-10;
const int MOD = int(1e9) + 7;
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;
}
template <class BidirectionalIterator>
bool next_partial_permutation(BidirectionalIterator first,
                              BidirectionalIterator middle,
                              BidirectionalIterator last) {
  reverse(middle, last);
  return next_permutation(first, last);
}
ll gcd(ll x, ll y) { return (x % y) ? gcd(y, x % y) : y; }
ll lcm(ll x, ll y) { return x / gcd(x, y) * y; }
ll GCD(VI v) {
  ll a = v[0];
  for (ll i = 1; i < SZ(v); i++) {
    a = gcd(a, v[i]);
  }
  return a;
}
ll LCM(VI v) {
  ll a = v[0];
  for (ll i = 1; i < SZ(v); i++) {
    a = lcm(a, v[i]);
  }
  return a;
}

struct LoopFinder {

  VI m_arr;
  VI m_vec;
  map<ll, ll> m_map;
  ll m_loop_until = 0;
  ll m_loop_length = 0;

  LoopFinder(VI &vec) { m_arr = vec; }

  void search_find() {
    // 初期条件
    ll s = 0;
    ll cnt = 1;
    m_vec.emplace_back(0);
    while (true) {
      s = m_arr[s];
      auto f = find(ALL(m_vec), s);
      if (f != m_vec.end()) {
        ll dst = distance(m_vec.begin(), f);
        m_loop_until = dst;
        m_loop_length = (cnt - m_loop_until);
        break;
      } else {
        m_vec.emplace_back(s);
      }
      cnt++;
    }
    // print(m_loop_until);
    // print(m_loop_length);
    // ここではm_loop_until=0,m_loop_length=3にってれば良い
    return;
  }

  ll main2(ll num_action) {
    search_find();
    ll ret = 0;
    // 移動回数がループに入る前に終わっていたら
    if (num_action < m_loop_until) {
      ret = m_vec[num_action];
    }
    // 移動回数がループに入ったあとの場合
    else {
      ll mod = (num_action - m_loop_until) % m_loop_length;
      ll index = mod + m_loop_until;
      ret = m_vec[index];
    }
    return ret + 1;
  }

  void search(ll s = 0) {
    // 初期条件
    m_map[0] = 0;
    ll cnt = 1;
    while (true) {
      ll index = m_arr[s];
      if (m_map[index]) {
        m_loop_until = m_map[index];
        m_loop_length = (cnt - m_loop_until);
        break;
      } else {
        m_map[index] = cnt;
      }
      cnt++;
    }
    print(m_loop_until);
    print(m_loop_length);
    return;
  }

  ll main(ll num_action) {
    search();
    ll ret = 0;
    if (num_action < m_loop_until) {
      ret = m_arr[num_action];
    } else {
      ll mod = (num_action - m_loop_until) % m_loop_length;
      ll index = mod + m_loop_until;
      ret = m_map[index];
    }
    return ret + 1;
  }
};

struct LoopFinder2 {
  VI vec;
  VI m_vec;
  VI m_vis;
  ll tmp;

  LoopFinder2(VI a, ll n) {
    vec = a;
    m_vis.resize(n, 0);
  }

  void search() {
    tmp = 0;
    while (!m_vis[tmp]) {
      m_vec.push_back(tmp);
      m_vis[tmp] = 1;
      tmp = vec[tmp];
    }
  }

  ll main(ll move) {
    search();
    ll ret = 0;
    rep(i, SZ(m_vec)) {
      if (m_vec[i] == tmp) {
        if (move <= i) {
          ret = m_vec[move] + 1;
          break;
        } else {
          ll len = SZ(m_vec) - i;
          move -= i;
          move %= len;
          ret = m_vec[i + move] + 1;
          break;
        }
      }
    }
    return ret;
  }
};

void loop_finder() {
  ll n;
  ll k;
  cin >> n >> k;
  vector<int> a(n);
  for (auto &x : a) {
    cin >> x;
    x--;
  }
  vector<int> vec, vis(n, 0);
  int v = 0;
  while (!vis[v]) {
    vec.push_back(v);
    vis[v] = 1;
    v = a[v];
  }
  for (int i = 0; i < vec.size(); i++) {
    if (vec[i] == v) {
      if (k <= i) {
        cout << vec[k] + 1 << '\\n';
      } else {
        ll len = vec.size() - i;
        k -= i;
        k %= len;
        cout << vec[i + k] + 1 << '\\n';
      }
    }
  }
}

void Main() {
  ll n, k;
  ll res = 0;
  cin >> n >> k;
  ll loop = 0;
  VI a(n);
  rep(i, n) cin >> a[i], --a[i]; // 2e5くらいまでならおｋ

  LoopFinder2 loopFinder(a, n);
  res = loopFinder.main(k);
  print(res);
  return;
}

int main() {
  std::cin.tie(0);
  std::ios_base::sync_with_stdio(false);
  std::cout << std::fixed << std::setprecision(15);
  Main();
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using ll = long long int;
using ld = long double;
using VI = vector<ll>;
using VD = vector<ld>;
using VVI = vector<VI>;
using VC = vector<char>;
using VVC = vector<VC>;
using VS = vector<string>;
using PLL = pair<ll, ll>;
using PLD = pair<ld, ld>;
using VPLL = vector<PLL>;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
#define REPD(i, n) for (ll i = (ll)(n)-1; i >= 0; i--)
#define ALL(x) (x).begin(), (x).end()
#define ALLR(x) (x).rbegin(), (x).rend()
#define SZ(x) ((ll)(x).size())
#define MAX(x) *max_element((x).begin(), (x).end())
#define MIN(x) *min_element((x).begin(), (x).end())
#define SORTR(x) sort((x).rbegin(), (x).rend())
#define SORT(x) sort((x).begin(), (x).end())
#define SUM(x) accumulate((x).begin(), (x).end(), 0)
#define FILL(x, a) fill(x.begin(), x.end(), a)
#define EACH(i, x)                                                             \\
  for (typeof((x).begin()) i = (x).begin(); i != (x).end(); ++i)
#define print(x) cout << x << "\\n"
#define EXIST(v, x) (std::find(v.begin(), v.end(), x) != v.end())

const ll INF = 1e18;
const ld EPS = 1e-10;
const int MOD = int(1e9) + 7;
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;
}
template <class BidirectionalIterator>
bool next_partial_permutation(BidirectionalIterator first,
                              BidirectionalIterator middle,
                              BidirectionalIterator last) {
  reverse(middle, last);
  return next_permutation(first, last);
}
ll gcd(ll x, ll y) { return (x % y) ? gcd(y, x % y) : y; }
ll lcm(ll x, ll y) { return x / gcd(x, y) * y; }
ll GCD(VI v) {
  ll a = v[0];
  for (ll i = 1; i < SZ(v); i++) {
    a = gcd(a, v[i]);
  }
  return a;
}
ll LCM(VI v) {
  ll a = v[0];
  for (ll i = 1; i < SZ(v); i++) {
    a = lcm(a, v[i]);
  }
  return a;
}

struct LoopFinder {
  VI m_vec;
  VI m_arr;
  map<ll, ll> m_map;
  ll m_loop_until = 0;
  ll m_loop_length = 0;

  LoopFinder(VI &vec) { m_vec = vec; }

  void search_find(ll s = 0) {
    ll ret = 0;
    ll cnt = 1;
    // 初期条件
    m_arr.emplace_back(0);
    while (true) {
      s = m_vec[s];
      // m_map[index]=cnt;
      auto f = find(ALL(m_arr), s);
      if (f != m_arr.end()) {
        ll dst = distance(m_arr.begin(), f);
        m_loop_until = dst;
        m_loop_length = (cnt - m_loop_until);
        //--m_loop_until;//はじめの１回分は省く
        break;
      } else {
        m_arr.emplace_back(s);
      }
      cnt++;
    }
    // print(m_loop_until);
    // print(m_loop_length);
    // ここではm_loop_until=0,m_loop_length=3にってれば良い
    return;
  }

  ll calculate2(ll num_action) {
    search_find();
    ll ret = 0;
    if (num_action < m_loop_until) {
      ret = m_arr[num_action];
    } else {
      ll mod = (num_action - m_loop_until) % m_loop_length;
      ll index = mod + m_loop_until;
      ret = m_arr[index];
    }
    return ret + 1;
  }

  void search(ll s = 0) {
    ll ret = 0;
    ll cnt = 0;
    while (true) {
      s = m_vec[s];
      // m_map[index]=cnt;
      if (m_map[s]) {
        m_loop_until = m_map[s];
        m_loop_length = (cnt - m_loop_until);
        return;
      } else {
        m_map[s] = cnt;
      }
      cnt++;
    }
    return;
  }

  ll calculate(ll num_action) {
    search();
    ll ret = 0;
    if (num_action < m_loop_until) {
      ret = m_vec[num_action];
    } else {
      ll mod = (num_action - m_loop_until) % m_loop_length;
      ll index = mod + m_loop_until;
      ret = m_map[index];
    }
    return ret + 1;
  }
};
/*
void loop_finder(){
        ll curr=0;
        bool b=true;
        for(ll i=1;i<=k;++i){
                b=false;
                ll tmp=abs(c[now]-i);
                i+=tmp(k-1)/tmp;
        }
        c[curr]=i;
        curr=a[curr];
        res=curr;
}*/

void Main() {
  ll n, k;
  ll res = 0;
  cin >> n >> k;
  ll loop = 0;
  VI a(n);
  rep(i, n) cin >> a[i], --a[i]; // 2e5くらいまでならおｋ

  LoopFinder loopFinder(a);
  res = loopFinder.calculate2(k);
  print(res);
  return;
}

int main() {
  std::cin.tie(0);
  std::ios_base::sync_with_stdio(false);
  std::cout << std::fixed << std::setprecision(15);
  Main();
}