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

#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

using namespace std;

int N;
long long int K;
int P[5050];
long long int C[5050];
int used[5050];

int main(void) {

  cin >> N >> K;
  for (int i = 1; i <= N; i++) {
    cin >> P[i];
  }
  for (int i = 1; i <= N; i++) {
    cin >> C[i];
  }

  long long int ans = -100000000000;

  for (int k = 1; k <= N; k++) {
    long long int sum = 0;
    long long int tmp = 0;
    int cnt = 0;
    for (int i = 1; i <= N; i++) {
      used[i] = 0;
    }
    int I = k;

    /*for (long long int i = 1; i <= K; i++) {
            I = P[I];
            tmp += C[I];
            ans = max(ans, tmp);
    }*/
    // cout << "I= " << k << " ans= " << ans << endl;

    if (K <= N) {
      for (int j = 1; j <= K; j++) {
        I = P[I];
        tmp += C[I];
        ans = max(ans, tmp);
      }
    } else {
      used[I] = 1;
      long long int tmp_K = K;
      while (1) {
        I = P[I];
        sum += C[I];
        tmp += C[I];
        cnt++;
        ans = max(ans, tmp);
        tmp_K--;
        if (used[I] == 1) {
          break;
        }
        used[I] = 1;
      }

      tmp = 0;
      for (int i = 1; i <= N; i++) {
        used[i] = 0;
      }
      used[I] = 1;
      tmp_K = K;
      while (1) {
        I = P[I];
        tmp += C[I];
        tmp_K--;
        long long int num = tmp_K / cnt;
        // cout << "tmp= " << tmp <<" num= " <<num<< endl;
        ans = max(ans, tmp + num * sum);

        if (used[I] == 1) {
          break;
        }
        used[I] = 1;
      }
    }
  }

  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

using namespace std;

int N;
long long int K;
int P[5050];
long long int C[5050];
int used[5050];

int main(void) {

  cin >> N >> K;
  for (int i = 1; i <= N; i++) {
    cin >> P[i];
  }
  for (int i = 1; i <= N; i++) {
    cin >> C[i];
  }

  long long int ans = -100000000000;

  for (int k = 1; k <= N; k++) {
    long long int sum = 0;
    long long int tmp = 0;
    int cnt = 0;
    for (int i = 1; i <= N; i++) {
      used[i] = 0;
    }
    int I = k;

    for (long long int i = 1; i <= K; i++) {
      I = P[I];
      tmp += C[I];
      ans = max(ans, tmp);
    }
    // cout << "I= " << k << " ans= " << ans << endl;
    /*
    if (K <= N) {
            for (int j = 1; j <= K; j++) {
                    I = P[I];
                    tmp += C[I];
                    ans = max(ans, tmp);
            }
    }
    else {
            used[I] = 1;
            long long int tmp_K = K;
            while (1) {
                    I = P[I]; sum += C[I]; tmp += C[I]; cnt++;  ans = max(ans,
    tmp); tmp_K--; if (used[I] == 1) { break; } used[I] = 1;
            }

            //cout << "cnt= " << cnt << endl;

            //if (sum > 0) {
                    I = k;
                    long long int num = K / cnt;
                    tmp = num*sum; tmp_K = K - num*cnt;

                    for (long long int i = 1; i <= tmp_K; i++) {
                            I = P[I]; tmp += C[I]; ans = max(ans, tmp);
                    //}
            }
    }*/
  }

  cout << ans << endl;

  return 0;
}