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

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

lint n, k;
vector<lint> p, c;

vector<vector<pair<lint, lint>>> doubling;

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}

pair<lint, lint> doubling_calc(lint start, lint x) {
  lint ans = 0;
  bitset<33> b(x);
  for (int i = 0; i < 33; i++) {
    if (b[i] == 1) {
      ans += doubling[start][i].second;
      start = doubling[start][i].first;
    }
  }
  return {start, ans};
}

// startから1~x回やる時の最大
lint simulate(lint start, lint x) {
  lint ans = LLONG_MIN;
  lint cost = 0;
  for (int i = 0; i < x; i++) {
    cost += c[p[start]];
    start = p[start];
    chmax(ans, cost);
  }
  return ans;
}

int main() {
  cin >> n >> k;
  p.resize(n);
  c.resize(n);
  for (int i = 0; i < n; i++) {
    cin >> p[i];
    p[i]--;
  }
  for (int i = 0; i < n; i++) {
    cin >> c[i];
  }
  // iから2^j回移動
  doubling.resize(n, vector<pair<lint, lint>>(33, {0, 0}));
  for (int i = 0; i < n; i++) {
    doubling[i][0] = {p[i], c[p[i]]};
  }
  for (int i = 0; i < 32; i++) {
    for (int j = 0; j < n; j++) {
      auto next = doubling[doubling[j][i].first][i];
      doubling[j][i + 1].first = next.first;
      doubling[j][i + 1].second = doubling[j][i].second + next.second;
    }
  }
  lint ans = LLONG_MIN;
  if (k <= 2 * n) {
    for (int i = 0; i < n; i++) {
      chmax(ans, simulate(i, k));
    }
  } else {
    for (int i = 0; i < n; i++) {
      chmax(ans, simulate(i, n));
    }
    for (int i = 0; i < n; i++) {
      auto before = doubling_calc(i, k - n * 2);
      lint cand = before.second + simulate(before.first, 2 * n);
      chmax(ans, cand);
    }
  }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

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

lint n, k;
vector<lint> p, c;

vector<vector<pair<lint, lint>>> doubling;

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}

lint simulate(lint start, lint x) {
  lint ans = 0;
  bitset<33> b(x);
  for (int i = 0; i <= 32; i++) {
    if (b[i]) {
      ans += doubling[start][i].second;
      start = doubling[start][i].first;
    }
  }
  return ans;
}

int main() {
  cin >> n >> k;
  p.resize(n);
  c.resize(n);
  for (int i = 0; i < n; i++) {
    cin >> p[i];
    p[i]--;
  }
  for (int i = 0; i < n; i++) {
    cin >> c[i];
  }
  // iから2^j回移動
  doubling.resize(n, vector<pair<lint, lint>>(33, {0, 0}));
  for (int i = 0; i < n; i++) {
    doubling[i][0] = {p[i], c[p[i]]};
  }
  for (int i = 0; i < 32; i++) {
    for (int j = 0; j < n; j++) {
      auto next = doubling[doubling[j][i].first][i];
      doubling[j][i + 1].first = next.first;
      doubling[j][i + 1].second = doubling[j][i].second + next.second;
    }
  }
  if (k <= n) {
    lint ans = LLONG_MIN;
    for (int i = 0; i < n; i++) {
      for (int j = 1; j <= k; j++) {
        chmax(ans, simulate(i, j));
      }
    }
    cout << ans << endl;
  } else {
    lint ans = 0;
    for (int i = 0; i < n; i++) {
      for (int j = 1; j <= n; j++) {
        chmax(ans, simulate(i, j));
      }
    }
    for (int i = 0; i < n; i++) {
      for (int j = max(1LL, k - 2 * n); j <= k; j++) {
        chmax(ans, simulate(i, j));
      }
    }
    cout << ans << endl;
  }
}