[AtCoder] ABC 175 D – Moving Piece

2020年8月16日2020年12月14日

問題

方針

マス \( i \) を選択したとき、\( N \) 回以下の移動でマス \( i \) に戻ってくるので、\( K \) 回以下の移動で何周できるかを考えます。周回しない方が最適な場合もあることに注意します。マス \( i \) を含むサイクルの長さを \( l_i \) とすると、マス \( i \) から \( n  (1 \leq  n < l_i) \) 回移動したとき、周回できる回数は、

\[ \left \lfloor \dfrac{K – n}{l_i} \right \rfloor\]

となります。したがって、全てのマスを全探索します。

コード

#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main() {
    int N;
    ll K;
    cin >> N >> K;
    int P[N];
    ll C[N];
    for (int i = 0; i < N; i++) {
        cin >> P[i];
        P[i]--;
    }
    for (int i = 0; i < N; i++) {
        cin >> C[i];
    }

    ll ans = -100000000000000;

    for (int i = 0; i < N; i++) {
        vector<ll> v;
        int p = P[i];
        ll c_sum = C[p];
        ll cnt = 1;
        ans = max(ans, c_sum);
        while (true) {
            p = P[p];
            if (P[i] == p) break;
            cnt++;
            if (cnt > K) break;
            c_sum += C[p];
            ans = max(ans, c_sum);
        }
        ll sum = 0;
        p = P[i];
        if (cnt > K) continue;
        ans = max(ans, c_sum * (K / cnt));
        for (int i = 0; i < cnt; i++) {
            sum += C[p];
            p = P[p];
            ans = max(ans, sum + c_sum * ((K - i - 1)/ cnt));
        }
    }
    cout << ans << "\n";

    return 0;
}