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

/*
    Author: Nguyen Tan Bao
    Status:
    Idea:
*/

#include <bits/stdc++.h>
#define FI first
#define SE second
#define EPS 1e-9
#define ALL(a) a.begin(), a.end()
#define SZ(a) int((a).size())
#define MS(s, n) memset(s, n, sizeof(s))
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define FORE(i, a, b) for (int i = (a); i >= (b); i--)
#define FORALL(it, a)                                                          \\
  for (__typeof((a).begin()) it = (a).begin(); it != (a).end(); it++)
//__builtin_ffs(x) return 1 + index of least significant 1-bit of x
//__builtin_clz(x) return number of leading zeros of x
//__builtin_ctz(x) return number of trailing zeros of x

using namespace std;
using ll = long long;
using ld = double;
typedef pair<int, int> II;
typedef pair<II, int> III;
typedef complex<ld> cd;
typedef vector<cd> vcd;

const ll MODBASE = 1000000007LL;
const int MAXN = 5010;
const int MAXM = 1000;
const int MAXK = 16;
const int MAXQ = 200010;

int n, k, p[MAXN], c[MAXN];
vector<int> a;
bool visited[MAXN];

ll solve() {
  int n = SZ(a);
  FOR(i, n, 2 * n - 1) a.push_back(a[i - n]);
  // FOR(i,0,SZ(a)-1) cout << a[i] << ' '; cout << endl;

  ll res = -1e18;
  FOR(i, 0, n - 1) res = max(res, (ll)a[i]);

  ll sum = 0;
  FOR(i, 0, n - 1) sum += a[i];

  if (sum >= 0) {
    ll times = k / n;
    ll w = sum * times;
    ll K = k % n;
    ll Max = -1e18;
    FOR(i, 0, n - 1) {
      ll now = 0;
      FOR(j, i, i + K - 1) {
        now += a[j];
        Max = max(Max, now);
      }
    }
    w += Max;
    res = max(res, w);

    if (times) {
      times--;
      w = sum * times;
      K = n;
      Max = -1e18;
      FOR(i, 0, n - 1) {
        ll now = 0;
        FOR(j, i, i + K - 1) {
          now += a[j];
          Max = max(Max, now);
        }
      }
      w += Max;
      res = max(res, w);
    }
  } else {
    ll K = min(k, n);
    ll Max = -1e18;
    FOR(i, 0, n - 1) {
      ll now = 0;
      FOR(j, i, i + K - 1) {
        now += a[j];
        Max = max(Max, now);
      }
    }
    res = max(res, Max);
  }

  return res;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(nullptr);
  cin >> n >> k;
  FOR(i, 1, n) cin >> p[i];
  FOR(i, 1, n) cin >> c[i];
  MS(visited, false);
  ll res = -1e18;
  FOR(i, 1, n)
  if (!visited[i]) {
    a.clear();
    int now = i;
    a.push_back(c[now]);
    visited[i] = true;
    while (p[now] != i) {
      now = p[now];
      a.push_back(c[now]);
      visited[now] = true;
    }
    res = max(res, solve());
  }
  cout << res;
  return 0;
}


Error: WA

Fixed submission:

/*
    Author: Nguyen Tan Bao
    Status:
    Idea:
*/

#include <bits/stdc++.h>
#define FI first
#define SE second
#define EPS 1e-9
#define ALL(a) a.begin(), a.end()
#define SZ(a) int((a).size())
#define MS(s, n) memset(s, n, sizeof(s))
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define FORE(i, a, b) for (int i = (a); i >= (b); i--)
#define FORALL(it, a)                                                          \\
  for (__typeof((a).begin()) it = (a).begin(); it != (a).end(); it++)
//__builtin_ffs(x) return 1 + index of least significant 1-bit of x
//__builtin_clz(x) return number of leading zeros of x
//__builtin_ctz(x) return number of trailing zeros of x

using namespace std;
using ll = long long;
using ld = double;
typedef pair<int, int> II;
typedef pair<II, int> III;
typedef complex<ld> cd;
typedef vector<cd> vcd;

const ll MODBASE = 1000000007LL;
const int MAXN = 5010;
const int MAXM = 1000;
const int MAXK = 16;
const int MAXQ = 200010;

int n, k, p[MAXN], c[MAXN];
vector<int> a;
bool visited[MAXN];

ll solve() {
  int n = SZ(a);
  FOR(i, n, 2 * n - 1) a[i] = a[i - n];

  ll res = -1e18;
  FOR(i, 0, n - 1) res = max(res, (ll)a[i]);

  ll sum = 0;
  FOR(i, 0, n - 1) sum += a[i];

  if (sum >= 0) {
    ll times = k / SZ(a);
    ll w = sum * times;
    ll K = k % SZ(a);
    ll Max = -1e18;
    FOR(i, 0, n - 1) {
      ll now = 0;
      FOR(j, i, i + K - 1) {
        now += a[j];
        Max = max(Max, now);
      }
    }
    w += Max;
    res = max(res, w);

    if (times) {
      times--;
      w = sum * times;
      K = n;
      Max = -1e18;
      FOR(i, 0, n - 1) {
        ll now = 0;
        FOR(j, i, i + K - 1) {
          now += a[j];
          Max = max(Max, now);
        }
      }
      w += Max;
      res = max(res, w);
    }
  } else {
    ll K = min(k, n);
    ll Max = -1e18;
    FOR(i, 0, n - 1) {
      ll now = 0;
      FOR(j, i, i + K - 1) {
        now += a[j];
        Max = max(Max, now);
      }
    }
    res = max(res, Max);
  }

  return res;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(nullptr);
  cin >> n >> k;
  FOR(i, 1, n) cin >> p[i];
  FOR(i, 1, n) cin >> c[i];
  MS(visited, false);
  ll res = -1e18;
  FOR(i, 1, n)
  if (!visited[i]) {
    a.clear();
    int now = i;
    a.push_back(c[now]);
    while (p[now] != i) {
      now = p[now];
      a.push_back(c[now]);
    }
    res = max(res, solve());
  }
  cout << res;
  return 0;
}