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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
#define FOR(i, a, b) for (ll i = (ll)a; i <= (ll)b; i++)
#define DEC(i, a, b) for (ll i = (ll)a; i >= (ll)b; i--)
typedef pair<ll, ll> pi;
typedef pair<ll, pi> pii;
typedef pair<pi, pi> pipi;
#define f first
#define s second
typedef vector<ll> vi;
typedef vector<pi> vpi;
typedef vector<pii> vpii;
#define pb push_back
#define pf push_front
#define all(v) v.begin(), v.end()
#define disc(v)                                                                \\
  sort(all(v));                                                                \\
  v.resize(unique(all(v)) - v.begin());
#define INF (ll)1e9 + 100
#define LLINF (ll)1e18
#define fastio                                                                 \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
mt19937
    rng(chrono::steady_clock::now()
            .time_since_epoch()
            .count()); // can be used by calling rng() or shuffle(A, A+n, rng)
inline ll rand(ll x, ll y) {
  ++y;
  return (rng() % (y - x)) + x;
} // inclusivesss

ll n, k, p[5005], c[5005], ans = -LLINF;

ll f(ll x) {
  ll res = -LLINF, cur = 0, st = x, num = 0;
  while (num < k) {
    num++;
    st = p[st];
    cur += c[st];
    res = max(res, cur);
    if (st == x)
      break;
  }
  if (num == k)
    return res;
  cur *= (k / num - 1);
  res = max(res, cur);
  FOR(i, 1, num + k % num) {
    st = p[st];
    cur += c[st];
    res = max(res, cur);
  }
  return res;
}

int main() {
  fastio;
  cin >> n >> k;
  FOR(i, 1, n) cin >> p[i];
  FOR(i, 1, n) cin >> c[i];
  FOR(i, 1, n) ans = max(ans, f(i));
  cout << ans;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
#define FOR(i, a, b) for (ll i = (ll)a; i <= (ll)b; i++)
#define DEC(i, a, b) for (ll i = (ll)a; i >= (ll)b; i--)
typedef pair<ll, ll> pi;
typedef pair<ll, pi> pii;
typedef pair<pi, pi> pipi;
#define f first
#define s second
typedef vector<ll> vi;
typedef vector<pi> vpi;
typedef vector<pii> vpii;
#define pb push_back
#define pf push_front
#define all(v) v.begin(), v.end()
#define disc(v)                                                                \\
  sort(all(v));                                                                \\
  v.resize(unique(all(v)) - v.begin());
#define INF (ll)1e9 + 100
#define LLINF (ll)1e18
#define fastio                                                                 \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
mt19937
    rng(chrono::steady_clock::now()
            .time_since_epoch()
            .count()); // can be used by calling rng() or shuffle(A, A+n, rng)
inline ll rand(ll x, ll y) {
  ++y;
  return (rng() % (y - x)) + x;
} // inclusivesss

ll n, k, p[5005], c[5005], ans = -LLINF;

ll f(ll x) {
  ll res = -LLINF, cur = 0;
  FOR(i, 1, k) {
    x = p[x];
    cur += c[x];
    res = max(res, cur);
  }
  return res;
}

int main() {
  fastio;
  cin >> n >> k;
  FOR(i, 1, n) cin >> p[i];
  FOR(i, 1, n) cin >> c[i];
  FOR(i, 1, n) ans = max(ans, f(i));
  cout << ans;
}