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

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

// 計算量 : O(n**2)
//   要素数が n で、1周期の頂点数が n の場合
int main() {
  int n, k;
  cin >> n >> k;
  vector<int> p(n), c(n);

  for (int i = 0; i < n; i++) {
    cin >> p.at(i);
  }
  for (int i = 0; i < n; i++) {
    cin >> c.at(i);
  }
  for (int i = 0; i < n; i++) {
    p.at(i)--;
  } // 0-indexed に変更

  ll ans = -1e18;
  for (int si = 0; si < n; si++) {
    int x = si;    // 開始位置
    vector<int> s; // 1周期の1つ1つのスコアを格納する
    ll total = 0;  // 1周期の合計のスコアを格納する

    // 周期を求めて vector s に詰める
    //
    // 入力例1の場合
    // x=0
    //   0 --> 1
    //   x=1 の値を加算
    // x=1
    //   1 --> 3
    //   x=3 の値を加算
    // x=3
    //   3 --> 0
    //   x=0 の値を加算
    //   x=si が true となるため、周期に入ったと判定しループを終了する
    //
    while (1) {
      x = p.at(x);
      s.push_back(c.at(x));
      total += c.at(x);

      // 周期に入ったら終了する
      if (x == si) {
        break;
      }
    }

    // 周期 s の中のどこで切るかを全探索する

    // 周期の頂点数
    int loop = s.size();
    // 累積和
    ll t = 0;

    // ### 周期の合計がプラスの場合
    //
    // 例
    //   k = 10
    //   p =  1  2  4 |  1  2  4 |  1  2  4 |  1
    //   c = +3 +4 -1 | +3 +4 -1 | +3 +4 -1 | +3
    //
    // i=0
    //   t = s[0] = 3
    //   now = t = 3
    //   周期の合計がプラスの場合
    //     e = (k-(i+1)) / loop = (10-(0+1)) / 3 = 9 / 3 = 3
    //     now = now + total + e = 3 + 6 * 3 = 21
    //   ans = max(-1e18, 21) = 21
    // i=1
    //   t = s[0] + s[1] = 3 + 4 = 7
    //   now = 7
    //   周期の合計がプラスの場合
    //     e = (k-(i+1)) / loop = (10-(1+1)) / 3 = 8 / 3 = 2
    //     now = now + total + e = 7 + 6 * 2 = 19
    // i=2
    //   t = s[0] + s[1] + s[2] = 3 + 4 - 1 = 6
    //   now = 6
    //   周期の合計がプラスの場合
    //     e = (k-(i+1)) / loop = (10-(2+1)) / 3 = 7 / 3 = 2
    //     now = now + total + e = 6 + 6 * 2 = 18
    //
    // si=0 の場合
    //   1コマ移動 + できるだけループ = 3 + (6*3) = 21
    //   2コマ移動 + できるだけループ = 7 + (6*2) = 19
    //   3コマ移動 + できるだけループ = 6 + (6*2) = 18
    //   --> これらより 21 が最大値となる
    // si=1 の場合
    //   1コマ移動 + できるだけループ = 4 + (6*3) = 22
    //   2コマ移動 + できるだけループ = 3 + (6*2) = 15
    //   3コマ移動 + できるだけループ = 6 + (6*2) = 18
    // si=2 の場合
    //   2 --> 4 --> 2 のループ
    //   省略
    // si=3 の場合
    //   1コマ移動 + できるだけループ = -1 + (6*3) = 17
    //   2コマ移動 + できるだけループ =  2 + (6*2) = 14
    //   3コマ移動 + できるだけループ =  6 + (6*2) = 18
    // si=4 の場合
    //   2 --> 4 --> 2 のループ
    //   省略
    //
    // 全スタート地点について「1回〜1周期の頂点数の回」だけ移動するパターンを全探索する。
    // 移動後、できるだけ周期でスコアを積み上げる。
    //
    //
    // ### 周期の合計がマイナスの場合
    //
    // ループするほどスコアが下がるので、ループしない。
    // 全スタート地点について、1回〜周期の頂点数の数だけ移動するパターンについて、
    // 合計値を全探索する。
    //
    for (int i = 0; i < loop; i++) {
      t += s.at(i);

      if (i + 1 > k) {
        break;
      }

      ll now = t;
      // 周期の合計がプラスになる場合
      if (total > 0) {
        // 最大移動回数から i+1 を引く
        // +1 しているのは、i が 0-indexed のため
        //   0番目の頂点に移動した際に、移動回数が1となる
        //   1番目の頂点に移動した際に、移動回数が2となる
        // 残りの移動回数をループ回数で割り、可能な限りループを行い、スコアを最大化させる
        ll e = (k - (i + 1)) / loop;
        now += total * e;
      }
      ans = max(ans, now);
    }
  }

  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

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

int main() {
  int n, k;
  cin >> n >> k;

  vector<ll> p(n + 1, 0);
  vector<ll> c(n + 1, 0);
  for (int i = 1; i <= n; i++) {
    cin >> p.at(i);
  }

  // 最低1回はコマを移動させる必要がある
  // 1回移動した際に取りうる最も低い点数も同時に算出する
  ll ans = 0;
  for (int i = 1; i <= n; i++) {
    cin >> c.at(i);
    ans = min(ans, c.at(i));
  }

  // n箇所スタート地点がある
  for (int i = 1; i <= n; ++i) {
    int next = p.at(i);
    ll cnt = 0;
    int move = k;
    while (move > 0) {
      ll score = c.at(next);
      cnt += score;
      move--;
      next = p.at(next);
      ans = max(ans, cnt);
    }
  }

  cout << ans << endl;
  return 0;
}