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

#include <algorithm>
#include <iostream>
#include <string>
#include <utility>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
#define REP(i, n) for (int i = 0; i < (n); i++)
#define RREP(i, n) for (int i = (n)-1; i >= 0; i--)
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define RFOR(i, a, b) for (int i = (a); i > (b); i--)
#define ALL(a) (a).begin(), (a).end()

ll N, K, P[5000], C[5000];

vector<vector<ll>> get_loops() {
  vector<vector<ll>> ret;
  bool checked[5000] = {};
  REP(i, N) {
    if (checked[i])
      continue;

    vector<ll> loop;
    ll index = i;
    while (true) {
      index = P[index];
      if (checked[index])
        break;
      loop.push_back(C[index]);
      checked[index] = true;
    }
    ret.push_back(loop);
  }
  return ret;
}

ll get_ans_sub(vector<ll> loop, ll len) {
  ll ret = 0, sum = 0, size = loop.size();
  REP(i, len) sum += loop[i];
  ret = sum;
  if (len == size) {
    return ret;
  }

  ll s = 0, e = len;
  REP(i, size) {
    sum += loop[e];
    sum -= loop[s];
    e = (e + 1) % size;
    s = (s + 1) % size;
    ret = max(ret, sum);
  }
  return ret;
}

ll get_ans(vector<ll> loop) {
  ll all_sum = 0;
  for (ll n : loop)
    all_sum += n;

  ll ret = -1e18, size = loop.size();
  FOR(i, 1, min(K, size) + 1) {
    ll cnt = (K - i) / size;
    ll ans_sub = get_ans_sub(loop, i) + cnt * max(0LL, all_sum);
    ret = max(ret, ans_sub);
  }
  return ret;
}

int main() {
  cin >> N >> K;
  REP(i, N) {
    cin >> P[i];
    P[i]--;
  }
  REP(i, N) cin >> C[i];

  vector<vector<ll>> loops = get_loops();
  ll ans = -1e18;
  for (auto loop : loops) {
    ans = max(ans, get_ans(loop));
  }
  cout << ans << endl;
  return 0;
}

Error: 0

Fixed submission:

#include <algorithm>
#include <iostream>
#include <string>
#include <utility>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
#define REP(i, n) for (int i = 0; i < (n); i++)
#define RREP(i, n) for (int i = (n)-1; i >= 0; i--)
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define RFOR(i, a, b) for (int i = (a); i > (b); i--)
#define ALL(a) (a).begin(), (a).end()

int N, K, P[5000], C[5000];

vector<vector<int>> get_loops() {
  vector<vector<int>> ret;
  bool checked[50000] = {};
  REP(i, N) {
    if (checked[i])
      continue;

    vector<int> loop;
    int index = i;
    while (true) {
      index = P[index] - 1;
      if (checked[index])
        break;
      loop.push_back(C[index]);
      checked[index] = true;
    }
    ret.push_back(loop);
  }
  return ret;
}

ll get_ans_sub(vector<int> loop, int len) {
  ll ret = 0, sum = 0;
  REP(i, len) sum += loop[i];
  ret = sum;

  int s = 0, e = len, size = loop.size();
  REP(i, N) {
    sum += loop[e];
    sum -= loop[s];
    e = (e + 1) % size;
    s = (s + 1) % size;
    ret = max(ret, sum);
  }
  return ret;
}

ll get_ans(vector<int> loop) {
  ll all_sum = 0;
  for (int n : loop)
    all_sum += n;

  ll ret = -1e18;
  int len = K % loop.size();
  if (len == 0)
    len = K;
  FOR(i, 1, len + 1) { ret = max(ret, get_ans_sub(loop, i)); }

  if (all_sum > 0 && K > loop.size()) {
    ll cnt = K / loop.size();
    if (K % loop.size() == 0)
      cnt--;
    ret += cnt * all_sum;
  }
  return ret;
}

int main() {
  cin >> N >> K;
  REP(i, N) cin >> P[i];
  REP(i, N) cin >> C[i];

  vector<vector<int>> loops = get_loops();
  ll ans = -1e18;
  for (auto loop : loops)
    ans = max(ans, get_ans(loop));
  cout << ans << endl;
  return 0;
}