AtCoderで期待値の和とDPを見分ける:ABC154 DをPythonで解く

読了 約9分 たびすけ
AtCoder ABC154 Dの期待値と固定長ウィンドウを示す図

次に読む記事

関連するテーマの記事を、先に確認できます。

固定窓か状態遷移かを先に分ける

確率・期待値の問題では、各試行の結果を後続の状態へ渡すかどうかを最初に確認します。選んだ長さKの隣接区間にあるサイコロの出目を足すだけなら、各サイコロの期待値を作り、区間の出入りだけを更新して比較できます。ABC154 D Dice in Lineはこの形の問題です。

一方で、出目によって次の状態や停止判断、到達確率が変わるなら、固定した窓の合計だけでは追えません。読後には、各要素の期待値を足してよい問題か、確率や期待値を状態として更新する問題かを、出力対象と状態の変化から切り分けられるようにします。

1からPまでのサイコロの期待値

1からPまでの値を同じ確率で取るサイコロでは、出目の期待値は値の平均です。したがって、公式解説が示すように、期待値は(P + 1) / 2になります。たとえばP = 4なら、出目1、2、3、4の平均は2.5です。

この式で、サイコロそのものを確率分布として持ち続けず、合計へ寄与する一つの数へ置き換えられます。次の表は、公式サンプルではない小さい例P = [1, 3, 2, 4]に式を適用したものです。

位置そのサイコロのP期待値
111
232
321.5
442.5

P=[1,3,2,4], K=2で窓を動かす

同じ例で、隣接する2個を選ぶ窓を左から動かします。最初の窓は位置1と2なので、期待値1.0 + 2.0 = 3.0です。右へ一つずらすと、位置3の1.5が入り、位置1の1.0が出ます。したがって、更新後は3.0 + 1.5 - 1.0 = 3.5です。

もう一度ずらすと、位置4の2.5を足し、位置2の2.0を引きます。窓の合計は3.5 + 2.5 - 2.0 = 4.0になります。この例で得られる窓の合計は左から3.0、3.5、4.0で、最大値は4.0です。この例は公式サンプルではありません。

窓入る期待値出る期待値更新後の合計
位置1〜2なしなし3.0
位置2〜31.51.03.5
位置3〜42.52.04.0

窓を毎回最初から足し直すのではなく、前の窓との差分だけを使うと、すでに計算したK - 1個分を再利用できます。これが、固定長の区間を比較するときにスライディングウィンドウを選べる理由です。

期待値の線形性で分かることと分からないこと

ABC154 Dで求めるのは、選んだ隣接K個の出目を合計した値の期待値です。公式解説では、各サイコロの出目の期待値を求め、その合計を窓ごとに比較します。したがって、位置iの期待値をEiと書けば、窓の値はそのK個のEiの和として扱えます。

ただし、この加法で分かるのは選んだ窓の合計の期待値です。ある条件が起きる確率、結果の分布、出目の後にどの状態へ進むかは、この式だけでは求まりません。問題文が確率そのものを求めているのか、残り回数や位置を状態として更新するのかを読み、求める量に応じて方法を切り替えます。

状態、初期値、遷移の置き方を先に確認したい場合は、AtCoder DP入門:Frog 1で状態・初期値・遷移を決めるが前提の整理に使えます。

ABC154 Dの入力条件と公式サンプル

ABC154 Dの公式問題では、左から並ぶN個のサイコロがあり、i番目のサイコロは1からpiまでの整数を同じ確率で出します。隣接するK個を選び、独立に振った出目の合計の期待値を最大化します。

制約は 1 ≤ K ≤ N ≤ 200000、1 ≤ pi ≤ 1000 です。入力値は整数で、出力する最大期待値は想定解との差の絶対誤差または相対誤差が 10-6 以下なら正解とされます。各窓でK個を足し直す方法は、最悪でO(NK)時間かかります。そこで、前の窓の合計を再利用します。

公式サンプル1

入力

5 3
1 2 2 4 5

出力

7.000000000000

公式サンプル2

入力

4 1
6 6 6 6

出力

3.500000000000

公式サンプル3

入力

10 4
17 13 13 12 15 20 10 13 17 11

出力

32.000000000000

Python参考コードと不変条件

コードは、まず各piから期待値の配列expectedを作ります。window = sum(expected[:K])で最初のK個を足し、その値を最初の窓としてanswerへ保存します。

N, K = map(int, input().split())
P = list(map(int, input().split()))
expected = [(value + 1) / 2 for value in P]

window = sum(expected[:K])
answer = window
for right in range(K, N):
    window += expected[right] - expected[right - K]
    answer = max(answer, window)

print(answer)

rightが一つ進むと、expected[right]が新しい右端として入り、expected[right - K]が古い左端として出ます。更新後のwindowは常に現在の連続K個の期待値の合計です。その直後にanswer = max(answer, window)を実行するため、answerは調べ終えた全窓の最大値を保持します。

期待値の配列を一度作る処理、最初の窓の計算、各要素を一度ずつ入れて一度ずつ出す走査を合わせると、時間計算量はO(N)です。expectedがN個の値を持つため、追加メモリはO(N)です。公式解説が扱う累積和で区間を求める方法と、本文の差分更新は、前に計算した区間の情報を再利用する点が共通しています。区間和の別の実装を確認したい場合は、AtCoderの累積和入門:Typical90 010をPythonで解く区間和の考え方へ進めます。

公式サンプルと別の境界を分けて確認する

公式サンプルとは別に、窓の個数が端になる入力も確認します。K = 1では一つずつの期待値を比較し、入力3 1 / 1 2 10の最大値は5.5です。K = Nでは窓が一つだけになり、入力3 3 / 1 2 10の合計は1 + 1.5 + 5.5 = 8.0です。どちらも公式サンプルとは異なる境界の確認です。

確率DPと期待値DPへ切り替える条件

固定窓の計算を続けるか迷ったら、出力対象と、試行後に持ち越す情報を分けて読みます。条件だけで確率が決まるなら、場合分けで直接求められることがあります。残りターン数に応じて終了時の値と続行時の期待値を選ぶなら、残りターン数を状態にした期待値の再帰になります。盤面位置と回数ごとに到達確率が変わるなら、その二つを状態にして確率を遷移させます。

したがって、期待値の和は「合計の期待値を比較する」という対象には使えますが、確率イベントや状態の移動を自動で置き換えるものではありません。問題文に次の状態、停止判断、到達確率が現れたら、固定窓の式を広げる前に状態と遷移を定めます。

次の練習

  1. ABC242 A T-shirt(公式解説):順位がA位以内なら確率1、A+1位からB位ならC / (B - A)、それより下なら0と場合分けします。DPを使わず、条件から確率を直接選ぶ練習です。
  2. ABC266 E Throwing the Die(公式解説):残りターン数だけを状態にし、各出目Dについて終了時のDと継続時のf(N - 1)の大きい方を平均します。期待値を状態として再帰する練習です。
  3. ABC275 E Sugoroku 4(公式解説):回数と盤面位置を状態にして到達確率を遷移させ、998244353の法で出力します。公式解説のdp[i][j]は、i回後にマスjにいる確率です。