0. はじめに
第七回日本最強プログラマー学生選手権~Advance~(AtCoder Beginner Contest 464) の D問題「Celester」 の解法例です。
動的計画法(DP)を用いて解いています。
「天気を変更する・しない」「前日の天気によってボーナスが入る」といった条件を、どのようにDPの表に落とし込んでいくかを解説します。
1. 問題概要
- 問題リンク: ABC 464 D - Celester
ざっくり言うと?
$N$ 日間の元の天気予報(晴れ S / 雨 W)が与えられます。各日の天気をそのままにするか、あるいはコストを払って反転させるかを選べます。
このとき、「前日が雨で、当日が晴れ」 になったときだけ嬉しさ(ボーナス)が貰えます。最終的な「嬉しさの合計 - コストの合計」の最大値を求める問題です。
ボーナスが発生するかどうかの判定には「前日の天気」と「今日の天気」の2日分の情報しか必要ありません。
このように「直前の状態だけを覚えておけば、次の日の最適解が作れる」という性質には、動的計画法(DP) が向いています。
2. DPの設計と遷移表
DPテーブルの定義
今回は以下のような2次元配列を考えます。
扱いやすいよう、晴れを0、雨を1として定義します。
-
dp[0][i]: $i$ 日目の天気を 晴れ(0) にしたときの、$i$ 日目までの嬉しさの最大値 -
dp[1][i]: $i$ 日目の天気を 雨(1) にしたときの、$i$ 日目までの嬉しさの最大値
※途中で嬉しさがマイナスになる(ペナルティが先行する)可能性があるため、初期値は -float('inf')(マイナス無限大)にしておきます。
1日目(初期化)
1日目は、元々の天気通りなら嬉しさ 0、天気を反転させるなら -cost[0] になります。
2日目以降の遷移(条件表)
$i$ 日目の状態を決めるために、「前日の天気(2通り)」×「今日の天気を変えるか(2通り)」の計4パターンを毎回ループで全探索します。
例えば今日の元の天気が晴れ(0)だった場合のそれぞれの組み合わせにおける「嬉しさの変動(happy)」は以下の表のようになります。
前日の天気 (pre) |
今日の元の天気 (s[i]) |
天気を変えるか (change) |
変更後の今日の天気 (next_state) |
嬉しさの変動 (happy) |
備考 |
|---|---|---|---|---|---|
| 0: 晴れ | 0: 晴れ | 0: 変えない | 0: 晴れ | 0 |
変動なし |
| 0: 晴れ | 0: 晴れ | 1: 変える | 1: 雨 | -cost[i] |
変更コストのみ |
| 1: 雨 | 0: 晴れ | 0: 変えない | 0: 晴れ | + bonus[i-1] |
雨 $\to$ 晴れボーナス獲得! |
| 1: 雨 | 0: 晴れ | 1: 変える | 1: 雨 | -cost[i] |
変更コストのみ |
元の天気が「雨(1)」の場合も同様に change によって今日の天気が決まり、「pre == 1 かつ next_state == 0」 のときのみボーナスが入ります。
また変更後の今日の天気 (next_state)は
「今日の元の天気 (s[i]) + 天気を変えるか (change)」を2で割った余り
として算出できます。
求めた happy を使って、以下のように毎ステップ最大値を更新していきます。
$$dp[\text{next_state}][i] = \max(dp[\text{next_state}][i], dp[\text{pre}][i-1] + \text{happy})$$
3. ACコード(PyPy3)
提出してACしたソースコードです。コメントアウトで処理の流れを追えるようにしています。
T = int(input())
for _ in range(T):
n = int(input())
S = input()
# 処理しやすいように、S(晴れ)を0、R(雨)を1に変換した配列sを作る
s = []
for i in range(n):
if S[i] == "S":
s.append(0)
else:
s.append(1)
cost = list(map(int, input().split()))
bonus = list(map(int, input().split()))
# dp[天気のステート][日数] で初期化
# 途中の点数がマイナスになることがあるのでマイナス無限大が初期値
dp = [[-float("inf")] * n for i in range(2)]
# 1日目の嬉しさを設定
dp[s[0]][0] = 0 # 元の天気のまま
dp[(s[0] + 1) % 2][0] = -cost[0] # 天気を変えた場合
# 2日目以降の嬉しさを計算
for i in range(1, n):
for pre in range(2): # 1日前の天気 (0:晴れ, 1:雨)
for change in range(2): # 天気を変えるか (0:変えない, 1:変える)
happy = 0
next_state = (s[i] + change) % 2 # 変更を決めた後の今日の天気
happy = -cost[i] * change # 変える場合はコストを引く
if (
pre == 1 and next_state == 0
): # 「前日が雨(1)で、今日が晴れ(0)」ならボーナス加算!
happy += bonus[i - 1]
# 暫定の最大値を更新
dp[next_state][i] = max(dp[next_state][i], dp[pre][i - 1] + happy)
# 最終日(N-1日目)の「晴れ」と「雨」のどちらか大きい方が答え
print(max(dp[0][-1], dp[1][-1]))
4. 計算量
-
時間計算量: $O(N)$
各ステップで2 × 2 = 4パターンの遷移を試すだけなので、日数 $N$ に比例した時間で終わります。 -
空間計算量: $O(N)$
2 × NのDPテーブルを保持するため $O(N)$ です。
まとめ
コンテスト時間には間に合わず、終了後の提出となりました。
次回は時間内に実装を終えられるようにしたいです。