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

#include <bits/stdc++.h>
using namespace std;
#define ll long long int
#define ld long double
#define LL_MAX LLONG_MAX
#define ULL_MAX ULLONG_MAX
#define ull unsigned long long
#define IOS                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);                                                               \\
  cout.tie(NULL)
#define pb(x) push_back(x)
#define re() pop_back()
#define eb(x) emplace_back(x)
#define forn(i, n) for (int i = 0; i < int(n); ++i)
#define fore(i, l, r) for (int i = int(l); i < int(r); ++i)
#define all(a) a.begin(), a.end()
#define msort(a) sort(a.begin(), a.end())
#define pii pair<int, int>
#define sz(x) int(x.size())
#define mk make_pair
#define endl '\\n'

const int maxn = 5e3 + 5;
vector<int> p(maxn + 1, 0);
vector<int> c(maxn, 0);
// map<int,int> mp;
vector<int> vis(maxn, -1);
vector<pair<int, ll>> num;
void solve() {
  int n, k;
  cin >> n >> k;
  for (int i = 1; i < n + 1; i++) {
    cin >> p[i];
  }
  for (int i = 1; i < n + 1; i++) {
    cin >> c[i];
  }

  ll ans = LLONG_MIN;

  for (int i = 1; i < n + 1; i++) {
    for (int i = 1; i < n + 1; i++) {
      vis[i] = -1;
    }

    // mp.clear();
    num.clear();
    ll mx = LLONG_MIN;
    num.eb(mk(i, 0));
    int count = 0;
    ll mpans = 0;
    while (count < k) {
      pair<int, ll> last = num.back();
      int go = p[last.first];

      if (vis[go] != -1) {
        int preind = vis[go];
        ll sm = num[preind - 1].second;
        if (last.second - sm <= 0) {
          count = k;
          break;
        }
        int ind = num.size() - 1 - (preind - 1);
        int val = (k - count) / ind;
        if (val > 0)
          val--;
        count += ind * val;
        mpans += val * (last.second - sm);
        mx = max(mx, mpans);
        break;
      }

      mpans += c[go];
      mx = max(mx, mpans);
      vis[go] = num.size();
      num.eb(mk(go, mpans));
      count++;
    }

    while (count < k) {
      pair<int, ll> last = num.back();
      int go = p[last.first];
      mpans += c[go];
      mx = max(mx, mpans);
      num.eb(mk(go, mpans));
      count++;
    }

    ans = max(ans, mx);
  }

  cout << ans;
}

int main() {
  IOS;
  ll t = 1;
  while (t--) {
    solve();
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define ll long long int
#define ld long double
#define LL_MAX LLONG_MAX
#define ULL_MAX ULLONG_MAX
#define ull unsigned long long
#define IOS                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);                                                               \\
  cout.tie(NULL)
#define pb(x) push_back(x)
#define re() pop_back()
#define eb(x) emplace_back(x)
#define forn(i, n) for (int i = 0; i < int(n); ++i)
#define fore(i, l, r) for (int i = int(l); i < int(r); ++i)
#define all(a) a.begin(), a.end()
#define msort(a) sort(a.begin(), a.end())
#define pii pair<int, int>
#define sz(x) int(x.size())
#define mk make_pair
#define endl '\\n'

const int maxn = 5e3 + 5;
vector<int> p(maxn + 1, 0);
vector<int> c(maxn, 0);
map<int, int> mp;
vector<pair<int, ll>> num;
void solve() {
  int n, k;
  cin >> n >> k;
  for (int i = 1; i < n + 1; i++) {
    cin >> p[i];
  }
  for (int i = 1; i < n + 1; i++) {
    cin >> c[i];
  }

  ll ans = LLONG_MIN;

  for (int i = 1; i < n + 1; i++) {
    mp.clear();
    num.clear();
    ll mx = LLONG_MIN;
    num.eb(mk(i, 0));
    int count = 0;
    ll mpans = 0;
    while (count < k) {
      pair<int, ll> last = num.back();
      int go = p[last.first];

      if (mp.count(go)) {
        int preind = mp[go];
        ll sm = num[preind - 1].second;
        if (last.second - sm <= 0) {
          count = k;
          break;
        }
        int ind = num.size() - 1 - (preind - 1);
        int val = (k - count) / ind;
        if (val > 0)
          val--;
        count += ind * val;
        mpans += val * (last.second - sm);
        mx = max(mx, mpans);
        break;
      }

      mpans += c[go];
      mx = max(mx, mpans);
      mp[go] = num.size();
      num.eb(mk(go, mpans));
      count++;
    }

    while (count < k) {
      pair<int, ll> last = num.back();
      int go = p[last.first];
      mpans += c[go];
      mx = max(mx, mpans);
      num.eb(mk(go, mpans));
      count++;
    }

    ans = max(ans, mx);
  }

  cout << ans;
}

int main() {
  IOS;
  ll t = 1;
  while (t--) {
    solve();
  }
}