Problem p02585 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <bits/stdc++.h>
using namespace std;

const long long INF = 1e18;

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> P(N), C(N);
  for (int i = 0; i < N; i++) {
    cin >> P[i];
    P[i]--;
  }
  for (int i = 0; i < N; i++) {
    cin >> C[i];
  }
  vector<int> val_len(N);
  vector<vector<long long>> val(N, vector<long long>());
  for (int i = 0; i < N; i++) {
    int ind = P[i];
    long long sum_tmp = 0LL;
    long long sum_max_tmp = -INF;
    while (true) {
      sum_tmp += C[ind];
      sum_max_tmp = max(sum_max_tmp, sum_tmp);
      val[i].push_back(C[ind]);
      if (ind == i)
        break;
      ind = P[ind];
    }
  }

  for (int i = 0; i < N; i++) {
    val_len[i] = val[i].size();
  }

  vector<vector<long long>> ruiseki_2(N, vector<long long>());
  vector<vector<long long>> ruiseki_max_2(N, vector<long long>());
  for (int i = 0; i < N; i++) {
    int ind = P[i];
    long long sum_tmp = 0LL;
    long long sum_max_tmp = -INF;
    ruiseki_2[i].push_back(sum_tmp);
    ruiseki_max_2[i].push_back(sum_max_tmp);
    for (int j = 0; j < 2 * val_len[i]; j++) {
      sum_tmp += C[ind];
      sum_max_tmp = max(sum_max_tmp, sum_tmp);
      ruiseki_2[i].push_back(sum_tmp);
      ruiseki_max_2[i].push_back(sum_max_tmp);

      ind = P[ind];
    }
  }

  long long ans = -INF;
  for (int i = 0; i < N; i++) {
    for (int j = 1; j <= min(2 * val_len[i], K); j++) {
      ans = max(ans, ruiseki_max_2[i][j]);
    }
    if (ruiseki_2[i][val_len[i]] > 0) {
      long long cnt = K / val_len[i];
      if (cnt > 0) {
        long long rui_tmp = (cnt - 1) * ruiseki_2[i][val_len[i]];
        long long K_remainder = K - (cnt - 1) * val_len[i];
        for (int j = 0; j <= K_remainder; j++) {
          if (cnt == 1 && j == 0)
            continue;
          ans = max(ans, rui_tmp + ruiseki_max_2[i][j]);
        }
      }
    }
  }

  // for (int i=0;i<N;i++) {
  // 	cout << "i: " << i << endl;
  // 	for (auto v : ruiseki_max_2[i]) {
  // 		cout << v << " ";
  // 	}
  // 	cout << endl;
  // }
  cout << ans << endl;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

const long long INF = 1e18;

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> P(N), C(N);
  for (int i = 0; i < N; i++) {
    cin >> P[i];
    P[i]--;
  }
  for (int i = 0; i < N; i++) {
    cin >> C[i];
  }
  vector<int> val_len(N);
  vector<vector<long long>> val(N, vector<long long>());
  vector<vector<long long>> ruiseki(N, vector<long long>());
  vector<vector<long long>> ruiseki_max(N, vector<long long>());
  for (int i = 0; i < N; i++) {
    int ind = P[i];
    long long sum_tmp = 0LL;
    long long sum_max_tmp = -INF;
    ruiseki[i].push_back(sum_tmp);
    ruiseki_max[i].push_back(sum_max_tmp);
    while (true) {
      sum_tmp += C[ind];
      sum_max_tmp = max(sum_max_tmp, sum_tmp);
      val[i].push_back(C[ind]);
      ruiseki[i].push_back(sum_tmp);
      ruiseki_max[i].push_back(sum_max_tmp);
      if (ind == i)
        break;
      ind = P[ind];
    }
  }

  for (int i = 0; i < N; i++) {
    val_len[i] = val[i].size();
  }

  vector<vector<long long>> ruiseki_2(N, vector<long long>());
  vector<vector<long long>> ruiseki_max_2(N, vector<long long>());
  for (int i = 0; i < N; i++) {
    int ind = P[i];
    long long sum_tmp = 0LL;
    long long sum_max_tmp = -INF;
    ruiseki_2[i].push_back(sum_tmp);
    ruiseki_max_2[i].push_back(sum_max_tmp);
    for (int j = 0; j < 2 * val_len[i]; j++) {
      sum_tmp += C[ind];
      sum_max_tmp = max(sum_max_tmp, sum_tmp);
      ruiseki_2[i].push_back(sum_tmp);
      ruiseki_max_2[i].push_back(sum_max_tmp);

      ind = P[ind];
    }
  }

  long long ans = -INF;
  for (int i = 0; i < N; i++) {
    for (int j = 1; j <= min(2 * val_len[i], K); j++) {
      ans = max(ans, ruiseki_max_2[i][j]);
    }
    if (ruiseki[i][val_len[i]] > 0) {
      long long cnt = K / val_len[i];
      if (cnt > 0) {
        long long rui_tmp = (cnt - 1) * ruiseki[i][val_len[i]];
        long long K_remainder = K - (cnt - 1) * val_len[i];
        for (int j = 0; j <= K_remainder; j++) {
          if (cnt == 1 && j == 0)
            continue;
          ans = max(ans, rui_tmp + ruiseki_max_2[i][j]);
        }
      }
    }
  }

  // for (int i=0;i<N;i++) {
  // 	cout << "i: " << i << endl;
  // 	for (auto v : ruiseki_max_2[i]) {
  // 		cout << v << " ";
  // 	}
  // 	cout << endl;
  // }
  cout << ans << endl;
}