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

#include <bits/stdc++.h>
using namespace std;
#define ALL(x) x.begin(), x.end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define debug(v)                                                               \\
  cout << #v << ":";                                                           \\
  for (auto x : v) {                                                           \\
    cout << x << ' ';                                                          \\
  }                                                                            \\
  cout << endl;
#define mod 1000000007
using ll = long long;
const int INF = 1000000000;
const ll LINF = 1001002003004005006ll;
int dx[] = {1, 0, -1, 0}, dy[] = {0, 1, 0, -1};
// ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

struct IOSetup {
  IOSetup() {
    cin.tie(0);
    ios::sync_with_stdio(0);
    cout << fixed << setprecision(12);
  }
} iosetup;

template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  os << p.first << " " << p.second;
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++)
    os << v[i] << (i + 1 == (int)v.size() ? "" : " ");
  return os;
}

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &x : v)
    is >> x;
  return is;
}

// 0-indexed parmutation only
vector<vector<int>> cycle_partition(const vector<int> &p) {
  int n = p.size();
  vector<vector<int>> ret;
  vector<bool> check(n, false);
  rep(i, n) if (!check[p[i]]) {
    vector<int> v;
    int pos = p[i];
    v.push_back(i);
    check[i] = true;
    while (pos != i) {
      v.push_back(pos);
      check[pos] = true;
      pos = p[pos];
    }
    ret.push_back(v);
  }
  return ret;
}

signed main() {
  int n;
  ll k;
  cin >> n >> k;
  vector<int> p(n);
  vector<ll> c(n);
  cin >> p;
  cin >> c;

  if (*max_element(ALL(c)) <= 0) {
    cout << *max_element(ALL(c)) << endl;
    return 0;
  }

  rep(i, n) p[i]--;
  auto h = cycle_partition(p);
  vector<vector<ll>> g(h.size());
  rep(i, (int)h.size()) {
    for (auto &x : h[i]) {
      g[i].push_back(c[x]);
    }
  }

  // debug(p);
  // for(auto &w:g){
  //     debug(w);
  // }
  // for(auto &w:h){
  //     debug(w);
  // }

  auto f = [&](vector<ll> &a) {
    // debug(a);
    int m = (int)a.size();
    vector<ll> sum(m + 1);
    sum[0] = 0;
    rep(i, m) sum[i + 1] = sum[i] + a[i];
    ll ret = 0;
    if (k <= m) {
      rep(i, k + 1) chmax(ret, sum[i]);
      return ret;
    }
    rep(i, m + 1) chmax(ret, sum[i]);
    {
      ll tmp = 0;
      ll t = k / m;
      if (t > 5)
        t -= 5;
      else
        t = 0;
      ll rem = k - m * t;
      tmp += t * sum.back();
      for (int i = 0; i <= rem; i++) {
        ll hoge = i / m;
        ll huga = i % m;
        chmax(ret, tmp + hoge * sum.back() + sum[huga]);
      }
    }
    return ret;
  };

  ll res = *max_element(ALL(c));
  for (auto &w : g) {
    int m = (int)w.size();
    rep(st, m) {
      vector<ll> v;
      rep(i, m) v.push_back(w[(st + i) % m]);
      chmax(res, f(v));
    }
  }
  cout << res << endl;
  return 0;
}


Error: 0

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define ALL(x) x.begin(), x.end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define debug(v)                                                               \\
  cout << #v << ":";                                                           \\
  for (auto x : v) {                                                           \\
    cout << x << ' ';                                                          \\
  }                                                                            \\
  cout << endl;
#define mod 1000000007
using ll = long long;
const int INF = 1000000000;
const ll LINF = 1001002003004005006ll;
int dx[] = {1, 0, -1, 0}, dy[] = {0, 1, 0, -1};
// ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

struct IOSetup {
  IOSetup() {
    cin.tie(0);
    ios::sync_with_stdio(0);
    cout << fixed << setprecision(12);
  }
} iosetup;

template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  os << p.first << " " << p.second;
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++)
    os << v[i] << (i + 1 == (int)v.size() ? "" : " ");
  return os;
}

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &x : v)
    is >> x;
  return is;
}

// 0-indexed parmutation only
vector<vector<int>> cycle_partition(const vector<int> &p) {
  int n = p.size();
  vector<vector<int>> ret;
  vector<bool> check(n, false);
  rep(i, n) if (!check[p[i]]) {
    vector<int> v;
    int pos = p[i];
    v.push_back(i);
    check[i] = true;
    while (pos != i) {
      v.push_back(pos);
      check[pos] = true;
      pos = p[pos];
    }
    ret.push_back(v);
  }
  return ret;
}

signed main() {
  int n;
  ll k;
  cin >> n >> k;
  vector<int> p(n);
  vector<ll> c(n);
  cin >> p;
  cin >> c;

  if (*max_element(ALL(c)) <= 0) {
    cout << *max_element(ALL(c)) << endl;
    assert(false);
    return 0;
  }

  rep(i, n) p[i]--;
  auto h = cycle_partition(p);
  vector<vector<ll>> g(h.size());

  rep(i, (int)h.size()) {
    for (auto &x : h[i])
      g[i].push_back(c[x]);
  }

  // for(auto &w:g){
  //     debug(w);
  // }

  auto f = [&](vector<ll> &a) {
    // debug(a);
    int m = (int)a.size();
    vector<ll> sum(m + 1);
    sum[0] = 0;
    rep(i, m) sum[i + 1] = sum[i] + a[i];
    ll ret = 0;
    if (k <= m) {
      rep(i, k + 1) chmax(ret, sum[i]);
      return ret;
    }
    rep(i, m + 1) chmax(ret, sum[i]);
    if (sum.back() > 0) {
      ll res = 0;
      ll t = k / m;
      res += t * sum.back();
      ll rem = k % m;
      rep(i, rem + 1) chmax(ret, res + sum[i]);
    }
    return ret;
  };

  ll res = 0;
  for (auto &w : g) {
    int m = (int)w.size();
    rep(st, m) {
      vector<ll> v;
      for (int i = st, j = 0; j < m; i = (i + 1) % m, j++)
        v.push_back(w[i]);
      chmax(res, f(v));
    }
  }
  cout << res << endl;
  return 0;
}