Problem p02585 (C++) — buggy submission, verdict: Runtime Error

#pragma GCC optimize("O3")
#pragma comment(linker, "/stack:200000000")
#pragma GCC optimize("unroll-loops")

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define int long long
#define pb push_back
#define pf push_front
#define eb emplace_back
#define mp make_pair
#define all(v) (v).begin(), (v).end()
#define rall(v) (v).rbegin(), (v).rend()
#define f first
#define s second
#define sz(x) (int)x.size()
#define endl "\\n"
#define forn(i, n) for (int i = 0; i < n; ++i)
#define fore(i, l, r) for (int i = int(l); i <= int(r); ++i)
#define rep(i, begin, end) for (__typeof(end) i = (begin); i != (end); i++)
#define fill(a, value) memset(a, value, sizeof(a));
#define gcd(a, b) __gcd((a), (b))
#define watch1(x) cout << (x) << endl
#define watch2(x, y) cout << (x) << " " << (y) << endl
#define watch3(x, y, z) cout << (x) << " " << (y) << " " << (z) << endl
#define fastio                                                                 \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<pii> vpii;
typedef tree<int, null_type, less<int>, rb_tree_tag,
             tree_order_statistics_node_update>
    oset;

const int INF = 9e18;
const int mod = 1e9 + 7;
const int N = 5005;

int i, j, n, k, p[N], c[N];

void solve() {
  cin >> n >> k;
  for (i = 1; i <= n; i++) {
    cin >> p[i];
  }
  for (i = 1; i <= n; i++) {
    cin >> c[i];
  }
  int ans = -INF;
  for (i = 1; i <= n; i++) {
    vi score, vis(n + 1);
    int start = p[i];
    while (!vis[start]) {
      score.pb(c[start]);
      vis[start] = 1;
      start = p[start];
    }
    int sum = 0;
    for (j = 0; j < sz(score); j++)
      sum += score[j];
    if (sum > 0) {
      int times = k / sz(score), rem = k % sz(score);
      int best1 = times * sum, pref = 0, mx = 0;
      for (j = 0; j < rem; j++) {
        pref += score[j];
        mx = max(mx, pref);
      }
      best1 += mx;
      ans = max(ans, best1);
      if (times) {
        int best2 = (times - 1) * sum, pref = 0, mx = 0;
        for (j = 0; j < sz(score); j++) {
          pref += score[j];
          mx = max(mx, pref);
        }
        best2 += mx;
        ans = max(ans, best2);
      }
    } else {
      int pref = 0, mx = -INF;
      for (j = 0; j < min(sz(score), k); j++) {
        pref += score[j];
        mx = max(mx, pref);
      }
      ans = max(ans, mx);
    }
  }
  cout << ans;
}

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

Error: WA

Fixed submission:

#pragma GCC optimize("O3")
#pragma comment(linker, "/stack:200000000")
#pragma GCC optimize("unroll-loops")

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define int long long
#define pb push_back
#define pf push_front
#define eb emplace_back
#define mp make_pair
#define all(v) (v).begin(), (v).end()
#define rall(v) (v).rbegin(), (v).rend()
#define f first
#define s second
#define sz(x) (int)x.size()
#define endl "\\n"
#define forn(i, n) for (int i = 0; i < n; ++i)
#define fore(i, l, r) for (int i = int(l); i <= int(r); ++i)
#define rep(i, begin, end) for (__typeof(end) i = (begin); i != (end); i++)
#define fill(a, value) memset(a, value, sizeof(a));
#define gcd(a, b) __gcd((a), (b))
#define watch1(x) cout << (x) << endl
#define watch2(x, y) cout << (x) << " " << (y) << endl
#define watch3(x, y, z) cout << (x) << " " << (y) << " " << (z) << endl
#define fastio                                                                 \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<pii> vpii;
typedef tree<int, null_type, less<int>, rb_tree_tag,
             tree_order_statistics_node_update>
    oset;

const int INF = 9e18;
const int mod = 1e9 + 7;
const int N = 5005;

int i, j, n, k, p[N], vis[N], c[N];

void solve() {
  cin >> n >> k;
  for (i = 1; i <= n; i++) {
    cin >> p[i];
  }
  for (i = 1; i <= n; i++) {
    cin >> c[i];
  }
  int ans = -INF;
  for (i = 1; i <= n; i++) {
    if (vis[p[i]])
      continue;
    vi score;
    int start = p[i];
    while (!vis[start]) {
      vis[start] = 1;
      start = p[start];
      score.pb(c[start]);
    }
    sort(rall(score));
    int sum = 0, best = 0;
    for (j = 0; j < sz(score); j++)
      sum += score[j];
    if (sum >= 0) {
      int rounds = k / sz(score) - 1, rem = k % sz(score);
      best += max(0LL, sum * rounds);
      int mx = -INF, pref = 0, best1 = 0;
      for (j = 0; j < k; j++) {
        pref += score[j];
        mx = max(mx, pref);
      }
      best += mx;
      mx = -INF, pref = 0;
      for (j = 0; j < min(k, sz(score)); j++) {
        pref += score[j];
        mx = max(mx, pref);
      }
      best1 = mx;
      ans = max({ans, best1, best});
    } else {
      int mx = -INF, pref = 0;
      for (j = 0; j < min(k, sz(score)); j++) {
        pref += score[j];
        mx = max(mx, pref);
      }
      best += mx;
      ans = max(ans, best);
    }
  }
  cout << ans;
}

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