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

#include <bits/stdc++.h>
#define ll long long
using namespace std;
#define ll long long
#define dd double
#define fio                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);                                                               \\
  cout.tie(NULL);
#define mod (ll)(998244353)
#define sz(a) (ll) a.size()
#define all(a) a.begin(), a.end()
#define fr(i, a, b) for (ll i = a; i < b; i++)
#define frr(i, a, b) for (ll i = a - 1; i >= b; i--)
#define tc(t)                                                                  \\
  ll t;                                                                        \\
  cin >> t;                                                                    \\
  while (t--)
#define pb emplace_back
#define ee emplace
#define rr return 0
#define mp make_pair
#define pr pair<ll, ll>
#define ff first
#define ss second
#define pie 3.1415926535
#define inf LLONG_MAX
ll mult(ll a, ll b, ll p = mod) { return ((a % p) * (b % p)) % p; }
ll add(ll a, ll b, ll p = mod) { return (a % p + b % p) % p; }
ll fpow(ll x, ll y) {
  ll res = 1;
  x = x % mod;
  if (x == 0)
    return 0;
  while (y > 0) {
    if (y & 1LL)
      res = (res * x) % mod;
    y = y >> 1LL;
    x = (x * x) % mod;
  }
  return res;
}
ll inv(ll a, ll p = mod) { return fpow(a, p - 2); }
bool sa(const pair<ll, ll> &a, const pair<ll, ll> &b) {
  return (a.second < b.second);
}
bool fd(const pair<ll, ll> &a, const pair<ll, ll> &b) {
  return (a.first > b.first);
}
bool sd(const pair<ll, ll> &a, const pair<ll, ll> &b) {
  return (a.second > b.second);
}
ll dx[4] = {0, 0, 1, -1};
ll dy[4] = {1, -1, 0, 0};
bool valid(ll x, ll y, ll n, ll m) {
  if (x < 0 || y < 0) {
    return false;
  } else if (x >= n || y >= m) {
    return false;
  } else
    return true;
}
#define MAXN 1000001
ll spf[MAXN];
void sieve() {
  spf[1] = 1;
  for (ll i = 2; i < MAXN; i++)

    // marking smallest prime factor for every
    // number to be itself.
    spf[i] = i;

  // separately marking spf for every even
  // number as 2
  for (ll i = 4; i < MAXN; i += 2)
    spf[i] = 2;

  for (ll i = 3; i * i < MAXN; i++) {
    // checking if i is prime
    if (spf[i] == i) {
      // marking SPF for all numbers divisible by i
      for (ll j = i * i; j < MAXN; j += i)

        // marking spf[j] if it is not
        // previously marked
        if (spf[j] == j)
          spf[j] = i;
    }
  }
}

int main() {
#ifndef ONLINE_JUDGE
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
#endif
  ll n, k;
  cin >> n >> k;
  map<ll, ll> mp;
  for (ll i = 1; i <= n; i++) {
    ll x;
    cin >> x;
    mp[i] = x;
  }
  ll new_city = 1;
  vector<ll> loop_cities;
  bool tt = true;
  map<int, int> val;
  loop_cities.push_back(1);
  val[1]++;
  while (k--) {
    new_city = mp[new_city];
    val[new_city]++;
    if (val[new_city] == 2) {
      int tt = true;
      break;
    } else {
      loop_cities.push_back(new_city);
    }
  }

  ll i = std::find(loop_cities.begin(), loop_cities.end(), new_city) -
         loop_cities.begin();
  if (tt) {
    ll x = loop_cities[i + (k % (loop_cities.size() - i))];
    cout << x;
  } else
    cout << new_city;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define ll long long
using namespace std;
#define ll long long
#define dd double
#define fio                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);                                                               \\
  cout.tie(NULL);
#define mod (ll)(998244353)
#define sz(a) (ll) a.size()
#define all(a) a.begin(), a.end()
#define fr(i, a, b) for (ll i = a; i < b; i++)
#define frr(i, a, b) for (ll i = a - 1; i >= b; i--)
#define tc(t)                                                                  \\
  ll t;                                                                        \\
  cin >> t;                                                                    \\
  while (t--)
#define pb emplace_back
#define ee emplace
#define rr return 0
#define mp make_pair
#define pr pair<ll, ll>
#define ff first
#define ss second
#define pie 3.1415926535
#define inf LLONG_MAX
ll mult(ll a, ll b, ll p = mod) { return ((a % p) * (b % p)) % p; }
ll add(ll a, ll b, ll p = mod) { return (a % p + b % p) % p; }
ll fpow(ll x, ll y) {
  ll res = 1;
  x = x % mod;
  if (x == 0)
    return 0;
  while (y > 0) {
    if (y & 1LL)
      res = (res * x) % mod;
    y = y >> 1LL;
    x = (x * x) % mod;
  }
  return res;
}
ll inv(ll a, ll p = mod) { return fpow(a, p - 2); }
bool sa(const pair<ll, ll> &a, const pair<ll, ll> &b) {
  return (a.second < b.second);
}
bool fd(const pair<ll, ll> &a, const pair<ll, ll> &b) {
  return (a.first > b.first);
}
bool sd(const pair<ll, ll> &a, const pair<ll, ll> &b) {
  return (a.second > b.second);
}
ll dx[4] = {0, 0, 1, -1};
ll dy[4] = {1, -1, 0, 0};
bool valid(ll x, ll y, ll n, ll m) {
  if (x < 0 || y < 0) {
    return false;
  } else if (x >= n || y >= m) {
    return false;
  } else
    return true;
}
#define MAXN 1000001
ll spf[MAXN];
void sieve() {
  spf[1] = 1;
  for (ll i = 2; i < MAXN; i++)

    // marking smallest prime factor for every
    // number to be itself.
    spf[i] = i;

  // separately marking spf for every even
  // number as 2
  for (ll i = 4; i < MAXN; i += 2)
    spf[i] = 2;

  for (ll i = 3; i * i < MAXN; i++) {
    // checking if i is prime
    if (spf[i] == i) {
      // marking SPF for all numbers divisible by i
      for (ll j = i * i; j < MAXN; j += i)

        // marking spf[j] if it is not
        // previously marked
        if (spf[j] == j)
          spf[j] = i;
    }
  }
}

int main() {
#ifndef ONLINE_JUDGE
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
#endif
  ll n, k;
  cin >> n >> k;
  map<ll, ll> mp;
  for (ll i = 1; i <= n; i++) {
    ll x;
    cin >> x;
    mp[i] = x;
  }
  ll new_city = 1;
  vector<ll> loop_cities;
  bool tt = true;
  loop_cities.push_back(1);
  while (k--) {
    new_city = mp[new_city];
    if (std::find(loop_cities.begin(), loop_cities.end(), new_city) !=
        loop_cities.end()) {
      int tt = true;
      break;
    } else {
      loop_cities.push_back(new_city);
    }
  }
  // for (auto i : loop_cities) {
  // 	cout << i << " ";
  // }

  ll i = std::find(loop_cities.begin(), loop_cities.end(), new_city) -
         loop_cities.begin();
  if (tt) {
    ll x = loop_cities[i + (k % (loop_cities.size() - i))];
    cout << x;
  } else
    cout << new_city;
  return 0;
}