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

#include <bits/stdc++.h>
using namespace std;
#define REP(i, b, e) for (int i = (b); i < (e); ++i)
#define RREP(i, b, e) for (int i = (b)-1; i >= e; --i)
#define rep(i, e) for (int i = 0; i < (e); ++i)

constexpr int MOD = 1000000007;
constexpr int INF = 1LL << 30;
constexpr long long LLINF = 1LL << 62;
constexpr double EPS = 1e-9;

struct Initializer {
  Initializer() {
    cin.tie(0);
    ios::sync_with_stdio(0);
    cout << fixed << setprecision(15) << boolalpha;
  }
} initializer;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (auto it = vec.begin(); it != vec.end(); ++it) {
    if (it != vec.begin())
      os << ", ";
    os << *it;
  }
  return os << "]";
}

template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &pa) {
  return os << "(" << pa.first << ", " << pa.second << ")";
}

template <typename Tuple, size_t N> struct TuplePrinter {
  static void print(const Tuple &t) {
    TuplePrinter<Tuple, N - 1>::print(t);
    cout << ", " << get<N - 1>(t);
  }
};

template <typename Tuple> struct TuplePrinter<Tuple, 1> {
  static void print(const Tuple &t) { cout << get<0>(t); }
};

template <typename... T>
ostream &operator<<(ostream &os, const tuple<T...> &tup) {
  os << "(";
  TuplePrinter<decltype(tup), sizeof...(T)>::print(tup);
  return os << ")";
}

template <typename T> ostream &operator<<(ostream &os, const set<T> &se) {
  os << "{";
  for (auto it = se.begin(); it != se.end(); ++it) {
    if (it != se.begin())
      os << ", ";
    os << *it;
  }
  return os << "}";
}

template <typename T, typename U>
ostream &operator<<(ostream &os, const map<T, U> &ma) {
  os << "{";
  for (auto it = ma.begin(); it != ma.end(); ++it) {
    if (it != ma.begin())
      os << ", ";
    os << it->first << ": " << it->second;
  }
  return os << "}";
}

inline void print(void) { cout << '\\n'; }
template <class T> inline void print(const T &x) { cout << x << '\\n'; }
template <class T, class... U> inline void print(const T &x, const U &...y) {
  cout << x << " ";
  print(y...);
}

int main() {
  long long n, k;
  cin >> n >> k;
  vector<int> a(n + 1);
  rep(i, n) cin >> a[i + 1];

  vector<bool> visited(n + 1, false);
  int cur = 1;
  while (visited[cur] == false) {
    visited[cur] = true;
    cur = a[cur];
  }
  int loopstart = cur;
  int firstcnt = 0;
  cur = 1;
  while (cur != loopstart) {
    ++firstcnt;
    cur = a[cur];
  }
  int loopcnt = 0;
  cur = loopstart;
  do {
    ++loopcnt;
    cur = a[cur];
  } while (cur != loopstart);

  if (k <= (long long)firstcnt) {
    int ans = 1;
    rep(i, (int)k) ans = a[ans];
    print(ans);
  } else {
    long long rem = (k - firstcnt) % loopcnt;
    int ans = 1;
    rep(i, firstcnt) ans = a[ans];
    rep(i, rem) ans = a[ans];
    print(ans);
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define REP(i, b, e) for (int i = (b); i < (e); ++i)
#define RREP(i, b, e) for (int i = (b)-1; i >= e; --i)
#define rep(i, e) for (int i = 0; i < (e); ++i)

constexpr int MOD = 1000000007;
constexpr int INF = 1LL << 30;
constexpr long long LLINF = 1LL << 62;
constexpr double EPS = 1e-9;

struct Initializer {
  Initializer() {
    cin.tie(0);
    ios::sync_with_stdio(0);
    cout << fixed << setprecision(15) << boolalpha;
  }
} initializer;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (auto it = vec.begin(); it != vec.end(); ++it) {
    if (it != vec.begin())
      os << ", ";
    os << *it;
  }
  return os << "]";
}

template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &pa) {
  return os << "(" << pa.first << ", " << pa.second << ")";
}

template <typename Tuple, size_t N> struct TuplePrinter {
  static void print(const Tuple &t) {
    TuplePrinter<Tuple, N - 1>::print(t);
    cout << ", " << get<N - 1>(t);
  }
};

template <typename Tuple> struct TuplePrinter<Tuple, 1> {
  static void print(const Tuple &t) { cout << get<0>(t); }
};

template <typename... T>
ostream &operator<<(ostream &os, const tuple<T...> &tup) {
  os << "(";
  TuplePrinter<decltype(tup), sizeof...(T)>::print(tup);
  return os << ")";
}

template <typename T> ostream &operator<<(ostream &os, const set<T> &se) {
  os << "{";
  for (auto it = se.begin(); it != se.end(); ++it) {
    if (it != se.begin())
      os << ", ";
    os << *it;
  }
  return os << "}";
}

template <typename T, typename U>
ostream &operator<<(ostream &os, const map<T, U> &ma) {
  os << "{";
  for (auto it = ma.begin(); it != ma.end(); ++it) {
    if (it != ma.begin())
      os << ", ";
    os << it->first << ": " << it->second;
  }
  return os << "}";
}

inline void print(void) { cout << '\\n'; }
template <class T> inline void print(const T &x) { cout << x << '\\n'; }
template <class T, class... U> inline void print(const T &x, const U &...y) {
  cout << x << " ";
  print(y...);
}

int main() {
  long long n, k;
  cin >> n >> k;
  vector<int> a(n + 1);
  rep(i, n) cin >> a[i + 1];

  vector<bool> visited(n, false);
  int cur = 1;
  while (visited[cur] == false) {
    visited[cur] = true;
    cur = a[cur];
  }
  int loopstart = cur;
  int firstcnt = 0;
  cur = 1;
  while (cur != loopstart) {
    ++firstcnt;
    cur = a[cur];
  }
  int loopcnt = 0;
  cur = loopstart;
  do {
    ++loopcnt;
    cur = a[cur];
  } while (cur != loopstart);

  if (k <= (long long)firstcnt) {
    int ans = 1;
    rep(i, k) ans = a[ans];
    print(ans);
  } else {
    long long rem = (k - firstcnt) % loopcnt;
    int ans = 1;
    rep(i, firstcnt) ans = a[ans];
    rep(i, rem) ans = a[ans];
    print(ans);
  }

  return 0;
}