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

#include <bits/stdc++.h>
#include <iomanip>
#include <iostream>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long LL;
#define REP(i, n) for (int i = 0; i < (n); i++)
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)

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

  vector<int> loop_count(N);
  REP(i, N) {
    LL init = P[i];
    LL count = 1;
    LL index = P[i] - 1;
    while (true) {
      LL Pi = P[index];
      if (Pi == init) {
        break;
      }
      index = Pi - 1;
      count++;
    }
    loop_count[i] = count;
  }
  // priority_queue<LL> que;
  LL max_point = C[P[0] - 1];

  REP(i, N) {
    LL loop = loop_count[i];
    LL point = 0;
    LL index = P[i] - 1;
    if (loop > K) {
      for (LL j = 0; j < K; j++) {
        point += C[index];
        if (max_point < point)
          max_point = point;
        index = P[index] - 1;
      }
    } else {
      for (LL j = 0; j < loop; j++) {
        point += C[index];
        // que.push(point);
        index = P[index] - 1;
      }
      // if(point < 0){continue;}

      LL one_loop_point = point;
      LL max_count = K / loop;
      LL r_count = K % loop;

      point = (one_loop_point > 0) ? one_loop_point * (max_count - 1) : 0;
      for (LL j = 0; j < loop; j++) {
        point += C[index];
        if (max_point < point)
          max_point = point;
        index = P[index] - 1;
      }

      for (LL j = 0; j < r_count; j++) {
        point += C[index];
        if (max_point < point)
          max_point = point;
        // que.push(point);
        index = P[index] - 1;
      }
    }
  }

  cout << max_point << endl;

  return 0;
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>
#include <iomanip>
#include <iostream>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long LL;
#define REP(i, n) for (int i = 0; i < (n); i++)
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)

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

  vector<int> loop_count(N);
  REP(i, N) {
    LL init = P[i];
    LL count = 1;
    LL index = P[i] - 1;
    while (true) {
      LL Pi = P[index];
      if (Pi == init) {
        break;
      }
      index = Pi - 1;
      count++;
    }
    loop_count[i] = count;
  }
  priority_queue<LL> que;

  REP(i, N) {
    LL loop = loop_count[i];
    LL point = 0;
    LL index = P[i] - 1;
    if (loop > K) {
      for (LL j = 0; j < K; j++) {
        point += C[index];
        que.push(point);
        index = P[index] - 1;
      }
    } else {
      for (LL j = 0; j < loop; j++) {
        point += C[index];
        que.push(point);
        index = P[index] - 1;
      }
      // if(point < 0){continue;}

      LL one_loop_point = point;
      LL max_count = K / loop;
      LL r_count = K % loop;

      point = one_loop_point * (max_count - 1);
      for (LL j = 0; j < loop; j++) {
        point += C[index];
        que.push(point);
        index = P[index] - 1;
      }

      for (LL j = 0; j < r_count; j++) {
        point += C[index];
        que.push(point);
        index = P[index] - 1;
      }
    }
  }

  LL ans = que.top();
  cout << ans << endl;

  return 0;
}