けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 175 D - Moving Piece (水色, 400 点)

異様に難しかった!!

問題へのリンク

問題概要

 1, 2, \dots, N の順列が与えられる。

順列の中の i から P[ i ] へ移動するとき、C[ P[ i ] ] だけスコアが加算される。

出発点を自由に選んで  1 回以上  K 回以下の移動を行うとき、得られるスコアの最大値を求めよ。

制約

  •  2 \le N \le 5000
  •  1 \le K \le 10^{9}

考えたこと

この問題は、ABC 167 D - Teleporter によく似ている。

drken1215.hatenablog.com

Teleporter の記事でも書いたように、 K 回の操作を行った結果を求める系の問題では

  • 同じ処理の繰り返しになっている箇所を見抜く
  • ダブリングする

というパターンが多い気がする。

解法 (1):同じ処理の繰り返しになっている箇所を見抜く

先ほど Teleporter に似ていると言ったけど、決定的に違うところが一つある。それは Teleporter では、下図みたいに、サイクルとなっている箇所から「髭」が伸びていたりするけど、、、

でも今回は、与えられるものが「順列」なので、かならず「いくつかのサイクル」になる。一個一個のサイクルは下図のようになっていて、どこから出発しても必ずいつかそこに戻ってくるようになっている。

たとえば、P = (2, 3, 4, 1, 6, 7, 5) だと

  • 1 -> 2 -> 3 -> 4 -> ...
  • 5 -> 6 -> 7 -> ...

という 2 つのサイクルに分けられる感じ。よって各サイクルごとに考察していこう!

サイクル中の K 個の要素の総和の最大化

問題は以下のように言い換えられる


  •  M 個の整数  c_{1}, \dots, c_{M} が与えられる (循環している)
  • それらから連続する  K 個以下を選んで総和を最大化せよ

 K \le 10^{9} という制約がとても厄介。でも大まかに言って次のパターンしかない

  • サイクルを一周して得られるスコアの総和が正のとき:最初に何個かとったあと、ひたすら目一杯周回する
  • サイクルを一周して得られるスコアの総和が 0 以下のとき: M 個のうちの連続する区間 (長さ  K 以下) の総和の最大値が答え

後者の場合は比較的簡単で、まず  K \le M であったならば、 K = M としてよくて、その下で連続する長さ  K 以下の区間の総和の最大値を求めれば OK。累積和をとっておけば  O(M^{2} でできる (このとき、サイクル全体を二週した数列の累積和をとると楽)。

前者の場合は、「最初に何個かとる部分の長さ」で場合分けすることにする。

  • 余りの長さを  r としたとき
  • 数列から連続する  r 個選んだ総和の最大値を求めて
  • 残りは  (K - r) / M 回周回する

というのが答えになる。注意点として、 r = 0 かつ  (K - r) / M = 0 であるような場合は除外する必要がある。

実装上はこれらを統一的に書ける!!!計算量は  O(N^{2})

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
int N;
long long K;
vector<int> P;
vector<long long> C;

long long solve() {
    long long res = -INF;

    // 順列を各サイクルに分解する
    vector<bool> used(N, false);
    vector<vector<long long>> ss;
    for (int i = 0; i < N; ++i) {
        if (used[i]) continue;
        int cur = i;
        vector<long long> s;
        while (!used[cur]) {
            used[cur] = true;
            s.push_back(C[cur]);
            cur = P[cur];
        }
        ss.push_back(s);
    }

    // 各サイクルごとに考える
    for (auto vec : ss) {
        long long M = vec.size();

        // サイクルを二週したものの累積和
        vector<long long> sum(M*2+1, 0);
        for (int i = 0; i < M*2; ++i) sum[i+1] = sum[i] + vec[i%M];

        // amari[r] := 連続する r 個の総和の最大値
        vector<long long> amari(M, -INF);
        for (int i = 0; i < M; ++i) {
            for (int j = 0; j < M; ++j) {
                chmax(amari[j], sum[i+j] - sum[i]);
            }
        }

        // 余りの長さで場合分け
        for (int r = 0; r < M; ++r) {
            if (r > K) continue;
            long long q = (K - r) / M;
            if (r == 0 && q == 0) continue;

            if (sum[M] > 0) chmax(res, amari[r] + sum[M] * q);
            else if (r > 0) chmax(res, amari[r]);
        }
    }
    return res;
}

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];
    cout << solve() << endl;
}

解法 (2):ダブリング

この手の問題におけるダブリング解法自体については以下の記事を参考に。

drken1215.hatenablog.com

ただし今回は、「ちょうど  K 回の総和を行う」のではなく、「 K 回以下の操作を行う」という設定なので、注意が必要!!!

  • next[ d ][ v ] := 頂点 v から  2^{d} だけ進んだ頂点
  • val[ d ][ v ] := 頂点 v から  2^{d} だけ進んだ中での総和
  • all[ d ][ v ] := 頂点 v から出発して  2^{d} 以下の距離進んだ中での総和の最大値

という感じのデータを管理すれば良さそう。

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
int N;
long long K;
vector<int> P;
vector<long long> C;

long long solve() {
    vector<vector<int>> next(60, vector<int>(N));
    vector<vector<long long>> val(60, vector<long long>(N)),
                              all(60, vector<long long>(N));
    for (int v = 0; v < N; ++v) {
        next[0][v] = P[v];
        val[0][v] = all[0][v] = C[v];
    }
    for (int d = 0; d + 1 < 60; ++d) {
        for (int v = 0; v < N; ++v) {
            next[d+1][v] = next[d][next[d][v]];
            val[d+1][v] = val[d][v] + val[d][next[d][v]];
            all[d+1][v] = max(all[d][v], val[d][v] + all[d][next[d][v]]);
        }
    }
    long long res = -INF;
    for (int v = 0; v < N; ++v) {
        long long sum = 0;
        long long offset = v;
        for (int d = 59; d >= 0; --d) {
            if (K & (1LL<<d)) {
                chmax(res, sum + all[d][offset]);
                sum += val[d][offset];
                offset = next[d][offset];
            }
        }
    }
    return res;
}

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];
    cout << solve() << endl;
}