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

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

const int N = 5e3 + 9;
int p[N], c[N], vis[N], ok;
long long ans, cc, dp[N][N], n, k;
unordered_map<int, int> mp, cost;
unordered_map<int, pair<int, int>> f;

const long long INF = -9999999999999999;

void func(int start, vector<int> &v) {
  if (mp[start] > 0) {
    ok = 1;
    return;
  }
  v.push_back(start);
  mp[start]++;
  func(p[start], v);
  return;
}

int main() {
  // freopen("input.txt", "r", stdin);
  ans = INF;
  // mp.clear();
  // cost.clear();
  cin >> n >> k;
  // 1-indexed based
  for (int i = 1; i <= n; i++) {
    cin >> p[i];
  }
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }
  /*int cycle = 1;
  memset(vis, 0, sizeof(vis));
  for (int i = 1; i <= n; i++) {
          if(!vis[i]) {
                  v.clear();
                  ok = 0;
                  cc = 0LL;
                  func(i, v);
                  reverse(v.begin(), v.end());
                  while(v[0] != v[v.size()-1])
                          v.pop_back();
                  reverse(v.begin(), v.end());
                  if(ok && v.size() > 1) {
                          for (int j = 0; j < v.size(); j++){
                                  cc += c[v[i]];
                                  f[v[i]] = make_pair(v[0], v.size());
  //starting point of the cycle If exists
                          }
                          //f[v[i]]
                          cost[v[0]] = cc;//contains the actual cost of the
  cycle
                  }
          }
  }*/

  long long ff, curmax, csum, cycleSum, t;
  vector<int> path;
  // k = k % (n+1);
  memset(dp, 0, sizeof(dp));
  for (int i = 1; i <= n; i++) {

    path.clear();
    ff = i;
    csum = 0LL;

    while (1) { // one cycle will always be there.
      // cout << "cycle\\n";
      csum += c[p[ff]];
      path.push_back(ff);
      ff = p[ff];
      if (ff == i)
        break;
    }
    // cout << "\\n";
    // for( int j = 0; j < (int) path.size(); j++) {
    //	cout << path[j] << " ";
    // }
    // cout << '\\n';
    cycleSum = 0LL;
    for (int j = 0; j < (int)path.size(); j++) {
      if (j + 1 > k) // cycle length is greater than 'k'
        break;
      cycleSum += c[path[j]];
      ans = max(ans, cycleSum);
    }
    // cycle length is smaller or sum of cycle is negative, so point in
    // extending this cycle
    if (cycleSum < 0 || (int)path.size() > k)
      continue;
    ans = max(ans, cycleSum);
    t = 0LL;
    long long other = 0LL;
    for (int j = 0; j < (int)path.size(); j++) {
      t += c[path[j]];
      if (j + 1 > k)
        break;
      other = t + ((k - j - 1) / (int)path.size()) *
                      csum; // number of times it can be wrapped around
      ans = max(ans, other);
    }
  }
  cout << ans << '\\n';
  return 0;
}

Error: TLE

Fixed submission:

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

const int N = 5e3 + 9;
int p[N], c[N], vis[N], ok;
long long ans, cc, dp[N][N], n, k;
unordered_map<int, int> mp, cost;
unordered_map<int, pair<int, int>> f;
const long long INF = -9999999999999999;
int main() {
  // freopen("input.txt", "r", stdin);
  ans = INF;
  // mp.clear();
  // cost.clear();
  cin >> n >> k;
  // 1-indexed based
  for (int i = 1; i <= n; i++) {
    cin >> p[i];
  }

  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }
  /*int cycle = 1;
  memset(vis, 0, sizeof(vis));
  for (int i = 1; i <= n; i++) {
          if(!vis[i]) {
                  v.clear();
                  ok = 0;
                  cc = 0LL;
                  func(i, v);
                  reverse(v.begin(), v.end());
                  while(v[0] != v[v.size()-1])
                          v.pop_back();
                  reverse(v.begin(), v.end());
                  if(ok && v.size() > 1) {
                          for (int j = 0; j < v.size(); j++){
                                  cc += c[v[i]];
                                  f[v[i]] = make_pair(v[0], v.size());
  //starting point of the cycle If exists
                          }
                          //f[v[i]]
                          cost[v[0]] = cc;//contains the actual cost of the
  cycle
                  }
          }
  }*/
  long long ff, curmax;
  // k = k % (n+1);
  memset(dp, 0, sizeof(dp));

  for (int i = 1; i <= n; i++) {
    cc = 0LL;
    ff = i;
    curmax = c[p[ff]];
    for (int j = 1; j <= k; j++) {
      cc += c[p[ff]];
      curmax = max(curmax, cc);
      ff = p[ff];
    }
    ans = max(ans, max(cc, curmax));
  }
  cout << ans << '\\n';
  return 0;
}