Problem p02684 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define rrep(i, n) for (int i = 1; i < (n); ++i)
#define all(a) (a).begin(), (a).end()
#define rall(a) (a).rbegin(), (a).rend()
#define maxs(a, b) a = max(a, b)
#define mins(a, b) a = min(a, b)
using namespace std;
typedef long long ll;
typedef pair<int, int> P;
const ll linf = 1001002003004005006ll;
const int inf = 1001001001;
const int mod = 1000000007;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n;
  ll k;
  cin >> n >> k;
  vector<int> a(n);
  rep(i, n) {
    cin >> a[i];
    a[i]--;
  }
  vector<int> def(n, 0);
  int res = 0;
  def[0] = 0;
  int cnt = 0;
  int rs = -1;
  for (ll i = 0; i < k; ++i) {
    cnt++;
    res = a[res];
    if (def[res]) {
      k -= def[res];
      rs = cnt - def[res];
      break;
    } else {
      def[res] = cnt;
    }
  }
  if (rs == -1) {
    cout << res + 1 << endl;
    return 0;
  }
  k %= rs;
  rep(i, k) { res = a[res]; }
  cout << res + 1 << endl;
  return 0;
}

Error: 0

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define rrep(i, n) for (int i = 1; i < (n); ++i)
#define all(a) (a).begin(), (a).end()
#define rall(a) (a).rbegin(), (a).rend()
#define maxs(a, b) a = max(a, b)
#define mins(a, b) a = min(a, b)
using namespace std;
typedef long long ll;
typedef pair<int, int> P;
const ll linf = 1001002003004005006ll;
const int inf = 1001001001;
const int mod = 1000000007;

class UnionFind {
public:
  // 親の番号を格納する。親だった場合は-(その集合のサイズ)
  vector<int> Parent;
  // 作るときはParentの値を全て-1にする
  // こうするとすべてバラバラになる
  UnionFind(int N) { Parent = vector<int>(N, -1); }
  // Aがどのグループに属しているか調べる
  int root(int A) {
    if (Parent[A] < 0)
      return A;
    return Parent[A] = root(Parent[A]);
  }
  // 自分のいるグループの頂点数を調べる
  int size(int A) {
    return -Parent[root(A)]; // 親をとってきたい
  }
  // AとBをくっつける
  bool connect(int A, int B) {
    // AとBを直接つなぐのではなく、root(A)にroot(B)をくっつける
    A = root(A);
    B = root(B);
    if (A == B) {
      // すでにくっついているからくっつけない
      return false;
    }
    // 大きい方(A)に小さいほう(B)をくっつけたい
    // 大小が逆だったらひっくり返しちゃう
    if (size(A) < size(B))
      swap(A, B);
    // Aのサイズを更新する
    Parent[A] += Parent[B];
    // Bの親をAに変更する
    Parent[B] = A;

    return true;
  }
  // AとBが同じグループにいるか
  bool same(int A, int B) { return root(A) == root(B); }
};

vector<bool> vi(200005);
vector<int> dist(200005);
vector<int> def[200005];
int rs;

void dfs(int s, int d = 0, int p = -1) {
  vi[s] = true;
  dist[s] = d;
  for (auto &c : def[s]) {
    if (c == p)
      continue;
    if (vi[c]) {
      rs = c;
      dist[c] = d - dist[c] + 1;
      continue;
    }
    vi[c] = true;
    dfs(c, d + 1, s);
  }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n;
  ll k;
  cin >> n >> k;
  UnionFind Uni(n);
  vector<int> a(n);
  rep(i, n) {
    cin >> a[i];
    a[i]--;
    Uni.connect(i, a[i]);
  }
  vector<int> res;
  rep(i, n) if (Uni.same(0, i)) res.emplace_back(i);
  rep(i, res.size()) { def[res[i]].emplace_back(a[res[i]]); }
  dfs(0);
  int ans = 0;
  int cnt = 0;
  while (1) {
    if (ans == rs) {
      k %= dist[rs];
      rep(i, k) { ans = a[ans]; }
      break;
    }
    ans = a[ans];
    k--;
    cnt++;
    if (cnt == k)
      break;
  }
  cout << ans + 1 << endl;
  return 0;
}