Problem p02585 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <complex>
#include <cstring>
#include <deque>
#include <exception>
#include <fstream>
#include <functional>
#include <iomanip>
#include <ios>
#include <iosfwd>
#include <iostream>
#include <istream>
#include <iterator>
#include <limits>
#include <list>
#include <locale>
#include <map>
#include <memory>
#include <new>
#include <numeric>
#include <ostream>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdexcept>
#include <streambuf>
#include <string>
#include <typeinfo>
#include <utility>
#include <valarray>
#include <vector>

using namespace std;
#define SZ(x) (int)(x.size())
#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i < b; ++i)
#define RREP(i, n) for (int i = n - 1; i >= 0; --i)
#define RFOR(i, a, b) for (int i = b - 1; i >= a; --i)

#define all(x) begin(x), end(x)
#define dump(x) cerr << #x << " = " << (x) << endl
#define debug(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

using ll = long long;
using vi = vector<int>;
using vvi = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;
using P = pair<int, int>;
const double eps = 1e-8;
const ll MOD = 1000000007;
const int INF = INT_MAX / 2;
const ll LINF = LLONG_MAX / 2;

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;
}

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

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  REP(i, SZ(v)) {
    if (i)
      os << ", ";
    os << v[i];
  }
  return os << "]";
}

template <class T1, class T2>
ostream &operator<<(ostream &os, const map<T1, T2> &mp) {
  os << "[";
  for (auto it = mp.begin(); it != mp.end(); ++it) {
    if (it != mp.begin()) {
      os << ", ";
    }
    os << *it;
  }
  return os << "]";
}

// edit
struct Solve {
  int N, K;
  vector<ll> cscore;
  vector<int> csize;
  vector<bool> visit;
  vector<int> P;
  vector<ll> C;

  void dfs(int v, int &sz, ll &score) {
    if (visit[v]) {
      // TODO
      return;
    }

    visit[v] = true;
    sz++;
    score += C[v];
    dfs(P[v], sz, score);

    csize[v] = sz;
    cscore[v] = score;
  }

  ll dfs2(int v, int sz) {
    if (sz == 0) {
      return 0ll;
    }
    ll ret = max(0ll, dfs2(P[v], sz - 1)) + C[P[v]];
    return ret;
  }

  void solve() {
    cin >> N >> K;
    cscore.resize(N, -1);
    csize.resize(N, -1);
    visit.resize(N, false);
    P.resize(N);
    C.resize(N);

    REP(i, N) cin >> P[i];
    REP(i, N) P[i]--;
    REP(i, N) cin >> C[i];

    REP(i, N) {
      int sz = 0;
      ll score = 0;
      if (!visit[i])
        dfs(i, sz, score);
    }

    ll ans = -LINF;

    for (int i = 0; i < N; ++i) {
      if (cscore[i] <= 0) {
        ll score = 0;
        ll add = dfs2(i, K);
        score += add;
        chmax(ans, score);
      } else {
        ll div = max((ll)K / csize[i] - 1, 0ll);
        ll score = div * cscore[i];
        ll nokori = K - div * csize[i];
        ll add = dfs2(i, nokori);
        score += add;
        chmax(ans, score);
      }
    }

    cout << ans << endl;
  }
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);
  //    std::ifstream in("input.txt");
  //    std::cin.rdbuf(in.rdbuf());

  Solve().solve();

  return 0;
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <complex>
#include <cstring>
#include <deque>
#include <exception>
#include <fstream>
#include <functional>
#include <iomanip>
#include <ios>
#include <iosfwd>
#include <iostream>
#include <istream>
#include <iterator>
#include <limits>
#include <list>
#include <locale>
#include <map>
#include <memory>
#include <new>
#include <numeric>
#include <ostream>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdexcept>
#include <streambuf>
#include <string>
#include <typeinfo>
#include <utility>
#include <valarray>
#include <vector>

using namespace std;
#define SZ(x) (int)(x.size())
#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i < b; ++i)
#define RREP(i, n) for (int i = n - 1; i >= 0; --i)
#define RFOR(i, a, b) for (int i = b - 1; i >= a; --i)

#define all(x) begin(x), end(x)
#define dump(x) cerr << #x << " = " << (x) << endl
#define debug(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

using ll = long long;
using vi = vector<int>;
using vvi = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;
using P = pair<int, int>;
const double eps = 1e-8;
const ll MOD = 1000000007;
const int INF = INT_MAX / 2;
const ll LINF = LLONG_MAX / 2;

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;
}

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

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  REP(i, SZ(v)) {
    if (i)
      os << ", ";
    os << v[i];
  }
  return os << "]";
}

template <class T1, class T2>
ostream &operator<<(ostream &os, const map<T1, T2> &mp) {
  os << "[";
  for (auto it = mp.begin(); it != mp.end(); ++it) {
    if (it != mp.begin()) {
      os << ", ";
    }
    os << *it;
  }
  return os << "]";
}

// edit
struct Solve {
  int N, K;
  vector<ll> cscore;
  vector<int> csize;
  vector<bool> visit;
  vector<int> P;
  vector<ll> C;

  void dfs(int v, int &sz, ll &score) {
    if (visit[v]) {
      // TODO
      return;
    }

    visit[v] = true;
    sz++;
    score += C[v];
    dfs(P[v], sz, score);

    csize[v] = sz;
    cscore[v] = score;
  }

  ll dfs2(int v, int sz) {
    if (sz == 0) {
      return 0ll;
    }
    ll ret = max(0ll, dfs2(P[v], sz - 1)) + C[P[v]];
    return ret;
  }

  void solve() {
    cin >> N >> K;
    cscore.resize(N, -1);
    csize.resize(N, -1);
    visit.resize(N, false);
    P.resize(N);
    C.resize(N);

    REP(i, N) cin >> P[i];
    REP(i, N) P[i]--;
    REP(i, N) cin >> C[i];

    REP(i, N) {
      int sz = 0;
      ll score = 0;
      if (!visit[i])
        dfs(i, sz, score);
    }

    ll ans = -LINF;

    for (int i = 0; i < N; ++i) {
      if (cscore[i] <= 0) {
        ll score = 0;
        ll add = dfs2(i, min(N, K));
        score += add;
        chmax(ans, score);
      } else {
        ll div = max((ll)K / csize[i] - 1, 0ll);
        ll score = div * cscore[i];
        ll nokori = K - div * csize[i];
        ll add = dfs2(i, nokori);
        score += add;
        chmax(ans, score);
      }
    }

    cout << ans << endl;
  }
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);
  //    std::ifstream in("input.txt");
  //    std::cin.rdbuf(in.rdbuf());

  Solve().solve();

  return 0;
}