期待値を計算する動的計画法(AtCoder ABC402E 問題)を解く

私は、動的計画法(DP)と期待値計算が結構苦手です。今回は、これら2つを組み合わせた問題(AtCoder 402E)があったので、解法を自分の忘却録として整理してみました。
はじめに
AtCoderの問題を後追いで解いていますが(コンテストには参加できてない)、動的計画法の遷移式の構築と、期待値計算というのがどうも苦手です。
今回は、動的計画法による期待値計算で解ける問題がAtCoder ABCコンテストのE問題として出題されていたので、整理ついでに記事にまとめました。

自分向けのメモとしての役割が強い記事ですが、参考になれば幸いです。
解法
ABC-402E問題
問題文については、リンク先を参考にしてください。
問題文のポイントを整理すると以下のようになります。
- 目的:
限られた所持金 $X$ 円を使って、正解した問題の得点の期待値の合計を最大化 - 条件:
- 問題は $N$ 問。
- 各問題 $i$に対し:
- 正解時の得点は $S_i$ 点。
- 提出には$C_i$ 円かかる。
- 正解する確率は $P_i%$(各回の提出は独立)。
- 所持金$X$ 円の範囲内で提出を行う。
- 同じ問題には複数回提出できるが、正解して得点を得られるのは1回のみ。
- 提出の結果を見てから、次にどの問題に提出するか選べる
- 求めるもの:
- 上記の条件で行動したときの、得点の期待値の最大値。
制約を確認すると、$N\le8$と少ないのがポイントです。それぞれの問題を解けた/解けてないを1bitで表すことにすると、8bit=256通りで列挙可能です。
コストも$X\le5000$と小さいです。
- $1\le N\le 8$
- $1\le S_i\le 2718$
- $1\le C_i\le X\le 5000$
- $1\le P_i\le 100$
- 入力される値は全て整数である
コスト(最大$X$)、状態(最大$2^N)を考えると動的計画法で解けるサイズに見えます。
状態遷移
自分で解いている時、状態遷移で結構悩みました。遷移としては以下の2つが考えられます。
- コスト$C_i$をかけて問題を解き、スコアを$S_i$を獲得
- コスト$C_i$をかけて問題を解いたが不正解でスコアを獲得なし
状態をdp[現在のコスト][現在の問題の集合]
とすると、それぞれ以下のような計算になります。
dp[curr-cost[i]][state+(1<<i)] + score[i]
dp[curr-cost[i]][state+(1<<i)]
は、i番目の問題を解いた状態のスコアで、これにscore[i]
を足したものdp[curr-cost[i]][state]
解いた問題は同じ状態だが、解く前よりコストがcost[i]
低い時の期待値
現在の期待値dp[curr][state]
は、この2つを$p_i$と$(1-p[i])$の比率で足し合わせたものにになります。

整理して考えると大した遷移ではないのですが、プログラムを考えている時には悩みまくりました。慣れの問題という人もいますが、この手の思考があまり向いてないのかも・・・・
状態の遷移式が決まったのであとは実装するだけです。
Pythonプログラム
以下は、問題を解くプログラムです。こうして整理して書くと、典型的な期待値DP + ビットDPの構成になっていることがわかります。
n,x = map(int,input().split())
dp = [[0]*(1<<n) for _ in range(x+1)]
score = []
cost = []
prob = []
for _ in range(n):
s,c,p = map(int,input().split())
score.append(s)
cost.append(c)
prob.append(p*0.01)
for curr in range(1,x+1):
for state in range(1<<n):
for i in range(n):
if (1 << i) & state: continue
if curr < cost[i]: continue
right = dp[curr-cost[i]][state+(1<<i)] + score[i]
wrong = dp[curr-cost[i]][state]
dp[curr][state] = max(dp[curr][state], right * prob[i] + wrong * (1-prob[i]) )
print(dp[x][0])
プログラムのループ部分を少し掘り下げます。
dp[curr][state]
とは何か?
これは、「現在の所持金が curr
円で、問題 state
(正解済の集合)にあるときの、最大期待得点」です。つまり、「今の状態から次に何かアクションをしたとき、どのような結果が得られるか?」を考える 現在の視点 となります。
right, wrongとは?
提出した結果、以下の2つのパターンに分岐します。
- 正解!
- 状態は
state + (1 << i)
(i番目も正解済みになる) - 所持金は
curr - cost[i]
- 得点は
score[i]
追加される dp[curr - cost[i]][state + (1 << i)]+score[i]
- 状態は
- 不正解
- 状態は変わらない(
state
のまま) - 所持金は
curr - cost[i]
- 得点は増えない
dp[curr - cost[i]][state]
- 状態は変わらない(
これらを確率付きで平均しすれば「提出による期待値」が得られます。
加重平均する
dp[curr][state]
には、「今この状態からどのような行動すれば最大得点になるか」が保存します。仮にi番目の問題を提出すれば、正解・不正解を加味するとどのような期待値になるかを計算するのが以下の行になります
dp[curr][state] = max(dp[curr][state], right * prob[i] + wrong * (1-prob[i]) )
最終的に、dp[X][0]
が答えになります。
まとめ
ビットDP(動的計画法)と期待値計算を組み合わせた問題を解いて、解説を行ってみました。こうやって解説しているとわかった気になるのですが、次に解けるかというと自信がありません。どうもこの手の問題は鬼門みたいです。

コンペ参加しているわけではないので、自分が解き始めるのと同時に生成AIにも解かせて競争しています。この手の典型問題は、生成AI(LLM)がサクッと解いてくるので少し自信がなくなりつつあります。類似問題がWebに大量に出回っている問題は、LLMがかなり解けるようになってきた印象です。