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

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

AtCoder ABC 306 D - Poisonous Full-Course (2Q, 茶色, 400 点)

シーケンシャルDP DP DP状態:on/off DP状態:state 最適化問題 最大スコア 解空間:O(2^N)通りの選択肢 典型要素を詰め合わせた教育的問題 NoviSteps2Q AtCoder AtCoder400点 ABC-D 茶色diff 負値が本質的に効く問題

とっても教育的な DP の問題！！！

問題概要

高橋君の前に N 個の料理が順に配られる。高橋君はその都度「食べる」「下げてもらう」を選択することができる。 i 番目の料理は、

  • X_{i} = 0 のとき、解毒剤入りの、美味しさが Y_{i} の料理
  • X_{i} = 1 のとき、毒入りの、美味しさが Y_{i} の料理

が配られる ( Y_{i} は負値もありうる)。

高橋君は、最初はお腹を壊していない。

  • お腹を壊していない状態で、解毒剤入りの料理を食べると、お腹を壊していない状態になる (変化しない)
  • お腹を壊していない状態で、毒入りの料理を食べると、お腹を壊してしまう
  • お腹を壊している状態で、解毒剤入りの料理を食べると、お腹を壊していない状態になる
  • お腹を壊している状態で、毒入りの料理を食べると、死亡してしまう

高橋君が死亡することなく、食べる料理の美味しさの総和を最大値を求めよ。

制約

  • 1 \le N \le 3 \times 10^{5}
  • -10^{9} \le Y_{i} \le 10^{9}

考えたこと

「解毒剤入りの料理はとにかく食べた方が良い」かと思いきや、美味しさが負の場合もあるなど、なかなかエゲツない設定になっている。Greedy では上手く行かなそうだ。こういうときは DP に限る！！！

今回の問題は選択肢が 2^{N} 通りあることに注意しよう (死亡しても食べ続けられるとした場合)。つまり、各料理に対して、「食べる」「下げてもらう」という 2 通りずつの選択肢があるので、全体で 2^{N} 通りの選択肢がある。このように、 2^{N} 通りの選択肢があるような問題では、DP が有効なことが多い。

さて、毎回の料理配布後の高橋君の状態として「お腹を壊していない」「お腹を壊している」が考えられる。このことから、次のような配列 dp を考えよう。

  • dp[i][0] ← 最初の i 個の料理を終えたあとに、高橋君がお腹を壊していない状態になっているときの、それまでに食べた料理の美味しさの総和の最大値
  • dp[i][1] ← 最初の i 個の料理を終えたあとに、高橋君がお腹を壊してる状態になっているときの、それまでに食べた料理の美味しさの総和の最大値

このとき、DP の遷移は次のように考えられる。まず、解毒剤が入っているかどうかにかかわらず「下げてもらう」という選択肢が常に存在する。「下げてもらう」という選択に対応する遷移をまず考えよう。

「下げてもらう」という選択に対応する DP 遷移

この場合は、高橋君のお腹の調子は変化しない。よって、次のように遷移を考えられる。

  • dp[i+1][0] = max(dp[i+1][0], dp[i][0])
  • dp[i+1][1] = max(dp[i+1][1], dp[i][1])

さらに、「食べる」場合も考えよう。

X_{i} = 0 (解毒剤入りのとき) のときに「食べる」とき

この場合は、「お腹を壊していない → そのまま」「お腹を壊している → 壊していない」の 2 パターンの遷移がある。

  • dp[i+1][0] = max(dp[i+1][0], dp[i][0] + Y[i])
  • dp[i+1][0] = max(dp[i+1][0], dp[i][1] + Y[i])

X_{i} = 1 (毒入りのとき) のときに「食べる」とき

この場合は、「お腹を壊していない → 壊す」という 1 パターンの遷移がある。

  • dp[i+1][1] = max(dp[i+1][1], dp[i][0] + Y[i])

コード

以上の DP を実装すれば AC となる。計算量は O(N) と評価できる。

#include <bits/stdc++.h>
using namespace std;
const long long INF = 1LL << 60;

template<class S, class T> inline bool chmax(S &a, T b) { return (a < b ? a = b, 1 : 0); }
template<class S, class T> inline bool chmin(S &a, T b) { return (a > b ? a = b, 1 : 0); }

int main() {
    int N;
    cin >> N;
    vector<long long> X(N), Y(N);
    for (int i = 0; i < N; ++i) cin >> X[i] >> Y[i];
    
    // 0: お腹を壊していない, 1: お腹を壊している
    vector dp(N + 1, vector(2, -INF));
    dp[0][0] = 0;
    for (int i = 0; i < N; ++i) {
        // 食べないとき: 0 -> 0, 1 -> 1
        chmax(dp[i+1][0], dp[i][0]);
        chmax(dp[i+1][1], dp[i][1]);
        
        // 食べるとき
        if (X[i] == 0) {
            // 解毒剤が効くので、0 -> 0, 1 -> 0
            chmax(dp[i+1][0], dp[i][0] + Y[i]);
            chmax(dp[i+1][0], dp[i][1] + Y[i]);
        } else {
            // 毒を受けるので、0 -> 1 の遷移のみ
            chmax(dp[i+1][1], dp[i][0] + Y[i]);
        }
    }
    cout << max(dp[N][0], dp[N][1]) << endl;
}