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

#include <bits/stdc++.h>
#define INF 1e9
using namespace std;

#define REPR(i, n) for (int i = (n); i >= 0; --i)
#define FOR(i, m, n) for (int i = (m); i < (n); ++i)
#define REP(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define ALL(a) (a).begin(), (a).end()
#define endl "\\n"

template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
typedef long long ll;

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

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

int turn = 0;
vector<int> poss; // 回数でどこにいたか
void dfs(const vector<int> &G, int v) {
  // cerr << "turn " << turn << " from: " << v << " to:" << G[v]<< endl;
  poss[turn] = v;
  turn++;
  seen[v] = true;
  hist.push(v);

  auto nv = G[v];

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

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

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

void solve() {
  ll N, K;
  cin >> N >> K;
  vector<int> g(N);
  REP(i, N) {
    int A;
    cin >> A;
    A--;
    g[i] = A;
  }

  // 探索
  seen.assign(N, false), finished.assign(N, false);
  poss.resize(N);
  dfs(g, 0);

  // サイクルを復元
  set<int> cycle;

  while (!hist.empty()) {
    int t = hist.top();
    cycle.insert(t);
    hist.pop();
    if (t == pos)
      break;
  }

  ll X;
  // compress
  if (hist.size() + cycle.size() <= K) {
    X = (K - hist.size()) % cycle.size() + hist.size();
  } else
    X = K;

  // cerr << "cycle size: " << cycle.size() << endl;
  // cerr << "hist size: " << hist.size() << endl;
  // cerr << "X: " << X << endl;
  cout << poss[X] + 1 << endl;
}

int main() {
  solve();
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define INF 1e9
using namespace std;

#define REPR(i, n) for (int i = (n); i >= 0; --i)
#define FOR(i, m, n) for (int i = (m); i < (n); ++i)
#define REP(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define ALL(a) (a).begin(), (a).end()
#define endl "\\n"

template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
typedef long long ll;

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

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

void dfs(const vector<int> &G, int v, int p) {
  seen[v] = true;
  hist.push(v);

  auto nv = G[v];

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

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

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

ll X;
int turn = 0;
int dfs2(const vector<int> &G, int p = 0) {
  turn++;
  // cerr << "turn " << turn << " from: " << p << " to:" << G[p] << endl;
  if (turn == X)
    return G[p];
  return dfs2(G, G[p]);
}

void solve() {
  ll N, K;
  cin >> N >> K;
  vector<int> g(N);
  REP(i, N) {
    int A;
    cin >> A;
    A--;
    g[i] = A;
  }

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

  // サイクルを復元
  set<int> cycle;

  while (!hist.empty()) {
    int t = hist.top();
    cycle.insert(t);
    hist.pop();
    if (t == pos)
      break;
  }

  // compress
  if (hist.size() + (2 * cycle.size()) < K) {
    X = (K - hist.size()) % cycle.size() + hist.size();
  } else
    X = K;

  // cerr << "cycle size: " << cycle.size() << endl;
  // cerr << "hist size: " << hist.size() << endl;
  // cerr << "X: " << X << endl;
  cout << dfs2(g) + 1 << endl;
}

int main() {
  solve();
  return 0;
}