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

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

int main() {
  int64_t N, K;
  cin >> N >> K;
  vector<int> P(N), C(N);
  for (int &e : P) {
    cin >> e;
    --e;
  }
  for (int &e : C)
    cin >> e;

  vector<int> done(N);
  int64_t ans = LLONG_MIN;
  for (int i = 0; i < N; ++i)
    if (!done[i]) {
      done[i] = 1;
      vector<int> cyc = {i};
      int at = i;
      int64_t ccost = C[i];
      while (P[at] != i) {
        at = P[at];
        cyc.push_back(at);
        ccost += C[at];
        done[at] = 1;
      }
      for (int a = 0; a < cyc.size(); ++a) {
        int64_t cost = 0;
        for (int b = (a + 1) % cyc.size(), mvs = 1;
             mvs <= min(K, (int64_t)cyc.size());
             b = (b + 1) % cyc.size(), ++mvs) {
          cost += C[cyc[b]];
          int rem = K - mvs;
          if (rem >= cyc.size() && ccost > 0) {
            ans = max(ans, cost + ccost * (rem / (int64_t)cyc.size()));
          } else {
            ans = max(ans, cost);
          }
        }
      }
    }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

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

int64_t sol(int64_t N, int64_t K, vector<int> P, vector<int> C) {
  vector<int> done(N);

  int64_t ans = LLONG_MIN;

  for (int i = 0; i < N; ++i) {
    if (done[i])
      continue;
    done[i] = 1;
    vector<int> cyc = {i};
    int64_t ccost = C[i];
    int at = i;
    while (P[at] != i) {
      at = P[at];
      ccost += C[at];
      cyc.push_back(at);
      done[at] = 1;
    }
    // cout << "here" << endl;
    // cout << "ee ";
    // for (int c : cyc) {
    // 	cout << c << " ";
    // }
    // cout << endl;

    int64_t mxextra = LLONG_MIN;
    for (int a = 0; a < cyc.size(); ++a) {
      int64_t extra = 0;
      int moves = 0;
      for (int b = (a + 1) % cyc.size();; b = (b + 1) % cyc.size()) {
        // cout << cyc.size() << " " << b << endl;
        ++moves;
        if (moves > K)
          break;
        extra += C[cyc[b]];
        mxextra = max(mxextra, extra);
      }
    }
    // cout << "here" << endl;

    // cout << mxextra << endl;

    ans = max(ans, mxextra);

    int64_t tmpans = 0;

    int rem = K;
    if (ccost > 0) {
      tmpans += (K / cyc.size()) * ccost;
      rem = K % cyc.size();
    }
    // cout << tmpans << " " << rem << endl;
    mxextra = 0;
    if (rem) {
      mxextra = LLONG_MIN;
      for (int a = 0; a < cyc.size(); ++a) {
        int64_t extra = 0;
        int moves = 0;
        for (int b = (a + 1) % cyc.size(); b != a; b = (b + 1) % cyc.size()) {
          ++moves;
          if (moves > rem)
            break;
          extra += C[cyc[b]];
          mxextra = max(mxextra, extra);
        }
      }
    }
    ans = max(ans, tmpans + mxextra);
  }
  return ans;
}

int64_t brute(int64_t N, int64_t K, vector<int> P, vector<int> C) {
  int64_t ans = LLONG_MIN;
  for (int i = 0; i < N; ++i) {
    int64_t cst = 0;
    int at = i;
    for (int j = 0; j < K; ++j) {
      cst += C[P[at]];
      at = P[at];
      ans = max(ans, cst);
    }
  }
  return ans;
}

int main() {
  // default_random_engine re;
  // for (int i = 0; i < 100000; ++i) {
  // 	int N = 3;
  // 	int K = re()%(N*2)+1;
  // 	vector<int> p(N), c(N);
  // 	for (int j = 0; j < N; ++j) {
  // 		p[j] = j;
  // 		c[j] = re()%(8)-4;
  // 	}
  // 	shuffle(p.begin(), p.end(), re);
  // 	cout << N << " " << K << endl;
  // 	for (int i = 0; i < N; ++i) {
  // 		cout << p[i]+1 << " ";
  // 	}
  // 	cout << endl;
  // 	for (int i = 0; i < N; ++i) {
  // 		cout << c[i] << " ";
  // 	}
  // 	cout << endl;
  // 	// sol(N,K,p,c);
  // 	cout << "== " << sol(N,K,p,c) << " " << brute(N,K,p,c) << endl;
  // 	assert(sol(N,K,p,c)==brute(N,K,p,c));
  // }
  // return 0;

  int64_t N, K;
  cin >> N >> K;
  vector<int> P(N), C(N);
  for (int &e : P) {
    cin >> e;
    --e;
  }
  for (int &e : C)
    cin >> e;

  cout << sol(N, K, P, C) << endl;
}