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

#include <algorithm>
#include <chrono>
#include <iostream>
#include <random>
#include <set>
#include <unordered_set>
#define rep(i, x) for (ll i = 0; i < x; i++)
using namespace std;
typedef long long ll;
const ll MOD = (ll)1e9 + 7;
const ll N = 5e5 + 5, inf = 5e18 + 5000;
ll add(ll x, ll y) {
  x += y;
  if (x >= MOD)
    return x - MOD;
  return x;
}
ll sub(ll x, ll y) {
  x -= y;
  if (x < 0)
    return x + MOD;
  return x;
}
ll mult(ll x, ll y) { return (x * y) % MOD; }
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  ll n, k;
  cin >> n >> k;
  vector<int> p(n);
  rep(i, n) {
    cin >> p[i];
    p[i]--;
  }
  vector<int> c(n);
  rep(i, n) cin >> c[i];

  vector<vector<int>> cycles(n, vector<int>());
  vector<ll> scores(n);
  for (int i = 0; i < n; i++) {
    unordered_set<int> seen;
    int curr = i;
    while (true) {
      if (seen.find(curr) != seen.end())
        break;
      seen.insert(curr);
      cycles[i].push_back(curr);
      scores[i] += c[curr];
      curr = p[curr];
    }
  }

  ll ans = -inf;
  for (int i = 0; i < n; i++) {

    int next = p[i];
    for (auto cyc_count :
         {0LL, k / (ll)cycles[next].size() - 1, k / (ll)cycles[next].size()}) {
      if (cyc_count < 0)
        continue;
      ll curr_score = 0;
      curr_score += cyc_count * scores[next];
      if (cyc_count != 0)
        ans = max(ans, curr_score);
      int curr = next;
      if (cyc_count != 0)
        ans = max(ans, curr_score);
      int last_block = cyc_count == k / (ll)cycles[next].size()
                           ? k % cycles[next].size()
                           : cycles[next].size();
      for (int j = 0; j < last_block && j < k; j++) {
        curr_score += c[curr];
        curr = p[curr];
        ans = max(ans, curr_score);
      }
    }
  }

  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <chrono>
#include <iostream>
#include <random>
#include <set>
#define rep(i, x) for (ll i = 0; i < x; i++)
using namespace std;
typedef long long ll;
const ll MOD = (ll)1e9 + 7;
const ll N = 5e5 + 5, inf = 5e18 + 5000;
ll add(ll x, ll y) {
  x += y;
  if (x >= MOD)
    return x - MOD;
  return x;
}
ll sub(ll x, ll y) {
  x -= y;
  if (x < 0)
    return x + MOD;
  return x;
}
ll mult(ll x, ll y) { return (x * y) % MOD; }
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  ll n, k;
  cin >> n >> k;
  vector<int> p(n);
  rep(i, n) {
    cin >> p[i];
    p[i]--;
  }
  vector<int> c(n);
  rep(i, n) cin >> c[i];

  vector<vector<int>> cycles(n, vector<int>());
  vector<ll> scores(n);
  for (int i = 0; i < n; i++) {
    set<int> seen;
    int curr = i;
    while (true) {
      if (seen.find(curr) != seen.end())
        break;
      seen.insert(curr);
      cycles[i].push_back(curr);
      scores[i] += c[curr];
      curr = p[curr];
    }
  }

  ll ans = -inf;
  for (int i = 0; i < n; i++) {

    int next = p[i];
    for (auto cyc_count :
         {0LL, k / (ll)cycles[next].size() - 1, k / (ll)cycles[next].size()}) {
      if (cyc_count < 0)
        continue;
      ll curr_score = 0;
      curr_score += cyc_count * scores[next];
      if (cyc_count != 0)
        ans = max(ans, curr_score);
      int curr = next;
      if (cyc_count != 0)
        ans = max(ans, curr_score);
      int last_block = cyc_count == k / (ll)cycles[next].size()
                           ? k % cycles[next].size()
                           : cycles[next].size();
      for (int j = 0; j < last_block && j < k; j++) {
        curr_score += c[curr];
        curr = p[curr];
        ans = max(ans, curr_score);
      }
    }
  }

  cout << ans << endl;

  return 0;
}