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

// i_m_arin

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef vector<long long> vll;
typedef vector<string> vs;
typedef priority_queue<ll> pqll;
typedef pair<ll, ll> pll;
typedef vector<pll> vpll;
typedef map<ll, ll> mpll;
typedef set<ll> stll;

#define f(i, a, n) for (ll i = a; i < (ll)n; i++)
#define f2(i, a, b) for (ll i = a; i <= (ll)b; i++)
#define pb push_back
#define ff first
#define ss second
#define pf push_front
#define mp make_pair
#define all(a) (a).begin(), (a).end()
#define decimal(x) fixed << setprecision(x)
#define fsort(a) sort(all(a))
#define rsort(a) sort(all(a), greater<>())
#define fast_io                                                                \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
#define mem0(a) memset(a, 0, sizeof(a))
#define mem1(a) memset(a, -1, sizeof(a))
#define lcm(a, b) (a * b) / __gcd(a, b)
#define ub(v, x) upper_bound(all(v), x) - v.begin()
#define lb(v, x) lower_bound(all(v), x) - v.begin()

const ld PI = acos(-1);
const ll MOD = 1e9 + 7;
const ll INF = LLONG_MAX;
const ll SIZE = 100000;

vll ar[1000000];
bool vis[1000000];
ll level[1000000];
ll pos = -1, pos2;

void dfs1(ll node, ll par) {
  if (par == 0)
    level[node] = 0;
  else
    level[node] = level[par] + 1;
  vis[node] = true;
  for (auto x : ar[node]) {
    if (!vis[x])
      dfs1(x, node);
  }
}

ll dfs(ll node, ll par) {
  vis[node] = true;
  for (auto x : ar[node]) {
    if (!vis[x]) {
      ll val = dfs(x, node);
      if (val != 0)
        return val;
    } else {
      if (x != par) {
        pos = x;
        pos2 = node;
        return abs(level[node] - level[x] + 1);
      }
    }
  }
  return 0;
}

// vll vect(1000000);

// void dfs2(ll cur,ll req){
// 	vis[cur]=true;
// 	vect.pb(cur);
// 	if(cur==req) return;
// 	for(auto x:ar[cur]){
// 		if(!vis[x]) dfs2(x,req);
// 	}
// }

void solve() {
  ll n, k;
  cin >> n >> k;
  ll a[n + 5];
  f(i, 0, 1000000) vis[i] = false;
  f(i, 1, n + 1) {
    ll x;
    cin >> x;
    a[i] = x;
    ar[i].pb(x);
  }
  dfs1(1, 0);

  f(i, 0, 1000000) vis[i] = false;
  ll ret = dfs(1, 0);
  if (pos == -1) {
    ll ans = 0, i = 1;
    while (n--) {
      ans = a[i];
      i = a[i];
    }
    cout << ans << endl;
  } else {
    ll ans = 0;

    ll fir = level[pos] - level[1];
    if (k <= fir + ret) {
      ll i = 1;
      while (k--) {
        ans = a[i];
        i = a[i];
      }
      cout << ans << endl;
    } else {

      k = k - (fir + ret);

      if (k > ret)
        k = k % ret;
      if (k == 0)
        k = ret;
      // f(i,0,1000000) vis[i]=false;
      // dfs2(pos,pos2);
      // vll vet(n+5);
      // f(i,0,vect.size()){
      // 	if(vect[i]>0){

      // 		vet[vect[i]]=a[vect[i]];
      // 	}
      // }

      ll i = pos;
      while (k--) {
        ans = a[i];
        i = a[i];
      }
      cout << ans << endl;
    }
  }
}

int main() {
  fast_io;
  int t = 1;
  // cin>>t;
  while (t--) {
    solve();
  }
  return 0;
}

Error: TLE

Fixed submission:

// i_m_arin

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef vector<long long> vll;
typedef vector<string> vs;
typedef priority_queue<ll> pqll;
typedef pair<ll, ll> pll;
typedef vector<pll> vpll;
typedef map<ll, ll> mpll;
typedef set<ll> stll;

#define f(i, a, n) for (ll i = a; i < (ll)n; i++)
#define f2(i, a, b) for (ll i = a; i <= (ll)b; i++)
#define pb push_back
#define ff first
#define ss second
#define pf push_front
#define mp make_pair
#define all(a) (a).begin(), (a).end()
#define decimal(x) fixed << setprecision(x)
#define fsort(a) sort(all(a))
#define rsort(a) sort(all(a), greater<>())
#define fast_io                                                                \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
#define mem0(a) memset(a, 0, sizeof(a))
#define mem1(a) memset(a, -1, sizeof(a))
#define lcm(a, b) (a * b) / __gcd(a, b)
#define ub(v, x) upper_bound(all(v), x) - v.begin()
#define lb(v, x) lower_bound(all(v), x) - v.begin()

const ld PI = acos(-1);
const ll MOD = 1e9 + 7;
const ll INF = LLONG_MAX;
const ll SIZE = 100000;

vll ar[1000000];
bool vis[1000000];
ll level[1000000];
ll pos = -1, pos2;

void dfs1(ll node, ll par) {
  if (par == 0)
    level[node] = 0;
  else
    level[node] = level[par] + 1;
  vis[node] = true;
  for (auto x : ar[node]) {
    if (!vis[x])
      dfs1(x, node);
  }
}

ll dfs(ll node, ll par) {
  vis[node] = true;
  for (auto x : ar[node]) {
    if (!vis[x]) {
      ll val = dfs(x, node);
      if (val != 0)
        return val;
    } else {
      if (x != par) {
        pos = x;
        pos2 = node;
        return abs(level[node] - level[x] + 1);
      }
    }
  }
  return 0;
}

// vll vect(1000000);

// void dfs2(ll cur,ll req){
// 	vis[cur]=true;
// 	vect.pb(cur);
// 	if(cur==req) return;
// 	for(auto x:ar[cur]){
// 		if(!vis[x]) dfs2(x,req);
// 	}
// }

void solve() {
  ll n, k;
  cin >> n >> k;
  ll a[n + 5];
  f(i, 0, 1000000) vis[i] = false;
  f(i, 1, n + 1) {
    ll x;
    cin >> x;
    a[i] = x;
    ar[i].pb(x);
  }
  dfs1(1, 0);

  f(i, 0, 1000000) vis[i] = false;
  ll ret = dfs(1, 0);
  if (pos == -1) {
    ll ans = 0, i = 1;
    while (k--) {
      ans = a[i];
      i = a[i];
    }
    cout << ans << endl;
  } else {
    ll ans = 0;

    ll fir = level[pos] - level[1];
    if (k <= fir + ret) {
      ll i = 1;
      while (k--) {
        ans = a[i];
        i = a[i];
      }
      cout << ans << endl;
    } else {

      k = k - (fir + ret);

      if (k > ret)
        k = k % ret;
      if (k == 0)
        k = ret;
      // f(i,0,1000000) vis[i]=false;
      // dfs2(pos,pos2);
      // vll vet(n+5);
      // f(i,0,vect.size()){
      // 	if(vect[i]>0){

      // 		vet[vect[i]]=a[vect[i]];
      // 	}
      // }

      ll i = pos;
      while (k--) {
        ans = a[i];
        i = a[i];
      }
      cout << ans << endl;
    }
  }
}

int main() {
  fast_io;
  int t = 1;
  // cin>>t;
  while (t--) {
    solve();
  }
  return 0;
}