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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MOD = 1000000007;
#define INF 1LL << 30
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(x) (x).begin(), (x).end()

using Graph = vector<vector<int>>;

// 探索
vector<bool> seen, finished;

// サイクル復元のための情報
int pos = -1;    // サイクル中に含まれる頂点 pos
stack<int> hist; // 訪問履歴

void dfs(const Graph &G, int v, int p) {
  seen[v] = true;
  hist.push(v);
  for (auto nv : G[v]) {
    // if (nv == p) continue; // 逆流を禁止する

    // 完全終了した頂点はスルー
    if (finished[nv])
      continue;

    // サイクルを検出
    if (seen[nv] && !finished[nv]) {
      pos = nv;
      return;
    }

    // 再帰的に探索
    dfs(G, nv, v);

    // サイクル検出したならば真っ直ぐに抜けていく
    if (pos != -1)
      return;
  }
  hist.pop();
  finished[v] = true;
}

int main() {
  // 頂点数 (サイクルを一つ含むグラフなので辺数は N で確定)
  ll k;
  int N;
  cin >> N >> k;

  int a[N];
  rep(i, N) {
    cin >> a[i];
    a[i]--;
  }

  Graph G(N);
  for (int i = 0; i < N; ++i) {
    G[i].push_back(a[i]);
  }

  // 探索
  seen.assign(N, false), finished.assign(N, false);
  pos = -1;
  dfs(G, 0, -1);

  // サイクルを復元
  set<int> cycle;
  while (!hist.empty()) {
    int t = hist.top();
    cycle.insert(t);
    hist.pop();
    if (t == pos)
      break;
  }
  // cout<<cycle.size()<<endl;
  int cnt = 1;
  int x = 0;
  while (1) {
    int y = a[x];
    if (cycle.count(x) && cycle.count(y)) {
      break;
    }
    if (cycle.count(x) == 0 && cycle.count(y)) {
      break;
    }
    cnt++;
    x = y;
  }
  ll jy = cnt + cycle.size() - 1;
  // cout<<cnt<<" "<<cycle.size<<endl;
  vector<int> ans;

  int num = jy;
  x = 0;
  rep(i, jy) {
    x = a[x];
    if (i >= cnt - 1)
      ans.push_back(x);
  }

  if (k <= jy) {
    x = 0;
    while (k > 0) {
      x = a[x];
      k--;
    }
    cout << x + 1 << endl;
  } else {
    k -= jy + 1;
    k %= ans.size();
    cout << ans[k] + 1 << endl;
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MOD = 1000000007;
#define INF 1LL << 30
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(x) (x).begin(), (x).end()

using Graph = vector<vector<int>>;

// 探索
vector<bool> seen, finished;

// サイクル復元のための情報
int pos = -1;    // サイクル中に含まれる頂点 pos
stack<int> hist; // 訪問履歴

void dfs(const Graph &G, int v, int p) {
  seen[v] = true;
  hist.push(v);
  for (auto nv : G[v]) {
    if (nv == p)
      continue; // 逆流を禁止する

    // 完全終了した頂点はスルー
    if (finished[nv])
      continue;

    // サイクルを検出
    if (seen[nv] && !finished[nv]) {
      pos = nv;
      return;
    }

    // 再帰的に探索
    dfs(G, nv, v);

    // サイクル検出したならば真っ直ぐに抜けていく
    if (pos != -1)
      return;
  }
  hist.pop();
  finished[v] = true;
}

int main() {
  // 頂点数 (サイクルを一つ含むグラフなので辺数は N で確定)
  ll k;
  int N;
  cin >> N >> k;

  int a[N];
  rep(i, N) {
    cin >> a[i];
    a[i]--;
  }

  Graph G(N);
  for (int i = 0; i < N; ++i) {
    G[i].push_back(a[i]);
  }

  // 探索
  seen.assign(N, false), finished.assign(N, false);
  pos = -1;
  dfs(G, 0, -1);

  // サイクルを復元
  set<int> cycle;
  while (!hist.empty()) {
    int t = hist.top();
    cycle.insert(t);
    hist.pop();
    if (t == pos)
      break;
  }

  int cnt = 1;
  int x = 0;
  while (1) {
    int y = a[x];
    if (cycle.count(x) && cycle.count(y)) {
      break;
    }
    if (cycle.count(x) == 0 && cycle.count(y)) {
      break;
    }
    cnt++;
    x = y;
  }
  ll jy = cnt + cycle.size() - 1;
  // cout<<cnt<<" "<<cycle.size<<endl;
  vector<int> ans;

  int num = jy;
  x = 0;
  rep(i, jy) {
    x = a[x];
    if (i >= cnt - 1)
      ans.push_back(x);
  }

  if (k <= jy) {
    x = 0;
    while (k > 0) {
      x = a[x];
      k--;
    }
    cout << x + 1 << endl;
  } else {
    k -= jy + 1;
    k %= ans.size();
    cout << ans[k] + 1 << endl;
  }
}