ABC402-E を Python で解いた

ABC402 E - Payment Required
https://atcoder.jp/contests/abc402/tasks/abc402_e

解説を読んでなんとか AC できました。期待値DPに慣れたいものです。

コンテスト中の自分の動き

A-D問題を解き終え、50分ほど残してE問題に取りかかりました。しかし最後まで入力例すら合わずに諦めました。自分の中では解法は見えていた気がしていて、どこがおかしいのか何故合わないのかと悩みました。後から振り返ってみるとかなり見当違いのことをしていたようです。残念な結果ではありますが、A-D問題の正答率が最近上がってきており安定感は増した気がします。

考察

貪欲法が頭をよぎりますが、$ N \leq 8 $ という制約から全探索できるのでは？と思いました。そして、$ 2 ^ 8 = 256 $、 $ X \leq 5000 $ ということから全ての状態を配列に持たせることができて動的計画法(DP)で解けるのではないかと思いました。

N 問の問題の回答状況を 0, 1 で表し、解けた問題を 1, 解けていない問題を 0 として 100101 のように表します。この例だと1, 4, 6問目が解けています。あとは DP をやっていくだけです。

しかし遷移で詰まりました。結局最後まで解くことができず、終わってから解説を読んでようやく理解できました。

DPの方針

問題の回答状況を bit, そのときの残りのお金を money とし、その状態で得られる総得点の期待値の最大値を $ DP[bit][money] $ に格納します。

状態遷移

https://atcoder.jp/contests/abc402/editorial/12715
公式解説にならい、現在正解できている問題を集合 T で表すことにします。

図で表すとこんな感じですね。正解すれば解けた問題が増え、残金が減ります。不正解の場合解けた問題は増えずに残金だけが減ります。

今、答えとして求めたい値は $ DP[0][X] $ です。このことから、0円からスタートしてゴールを X 円とし、そこまで上っていくイメージで解いていきます。つまり状態遷移の矢印の向きを逆にします。受け取る（もらう）DPになります。

これを式にすると
$ dp[T][x] = (dp[T∪{i}][x - C_i] + S[i]) * (P_i / 100) + dp[T][x - C_i] * ((100 - P_i) / 100) $
となります。

自分の過ち

私の場合は先に書いた「矢印を逆にする」発想が出てきませんでした。なので DP の開始点を DP[0][X] とし、答えを DP 配列の中で最大のものにしていました。初期値は -1, DP[0][X] だけ 0 にして……と、今から思えばよくわからないことをしていました。こんなコードです。

for money in range(5000, -1, -1):
    if dp[num][money] == -1: continue

    for i in range(N):
      if patterns[i] == 1: continue # もう解けている問題はスルー
      if money - C[i] >= 0:
        # 問題 i を解けたとき, お金が C 減って得点が S 増える。期待値にするので P / 100 をかける。
        dp[num + 2**(N-1-i)][money - C[i]] = \
          max(dp[num + 2**(N-1-i)][money - C[i]], dp[num][money] + S[i] * P[i] / 100)
                 
        # 問題 i を解けなかったとき
        dp[num][money - C[i]] = max(dp[num][money - C[i]], dp[num][money])

状態遷移で詰まりました。私がやろうとしたのは問題 i が解けたときは $ S_i \times \frac{P_i}{100} $ を足して、解けなかったときはそのままの値にするというものでした。確率 $ 1 - \frac{P_i}{100} $を使わないのは変では？という疑問がずっと消えませんでした。

• ゴールの設定を見誤った（しっかりと考えられていなかった）せいでDPの方向がおかしかった
• 【追記】そもそもDPテーブルの定義が曖昧だった。正解数0、残金X、得点0からスタートしてdp[T][x]を「その状態において今持っている得点の最大値の期待値」とするのか、正解数0、残金Xをゴールとしてdp[T][x]を「その状態からスタートしてお金を使い切るまでに得られる得点の最大値の期待値」とするのかは全然違う。この問題は後者であることにはっきりと気づかないといけない。「Tとxを元にDPテーブルを作る」程度のふんわりした見通ししか立っていないのにコードを書き始めてはいけない。
• 期待値の遷移の計算がわかっていなかった（前の期待値に次の得点を足してから確率$P$を掛けないといけない）
• ついでに、bit 計算に慣れておらず itertools を使ったり遠回りをしてしまっていた

あたりが自分の反省点、要改善点です。DPはテーブル設計だけでなく欲しい答えがどこにあるのかをしっかり考えてから書いた方がいいですね。

実装

ほとんど公式解説の変数名を変えただけのコードです。先の式でいう変数 T を bit に、変数 x を money にしています。

N, X = map(int, input().split())
S = [0 for _ in range(N)]
C = [0 for _ in range(N)]
P = [0 for _ in range(N)]
for i in range(N):
  S[i], C[i], P[i] = map(int, input().split())

dp = [[0.0 for _ in range(5001)] for _ in range(2**N)]
# 目的は dp[0][X] を求めること。

# 受け取る DP
for money in range(X+1): # 残金 money 円のとき
  for bit in range(1 << N): # 問題の回答状況が bit のとき
    for i in range(N):
      if money - C[i] < 0: continue
      if bit & (1<<i) == (1<<i): continue # 既に問題 i が解けている場合はスルー
      
      # 問題 i を解いた結果 dp[bit][money] になる場合
      solved_case = dp[bit | (1 << i)][money - C[i]]
      # 問題 i を解けなかった結果 dp[bit][money] になる場合
      unsolved_case = dp[bit][money - C[i]]

      dp[bit][money] = max(dp[bit][money], \
        (solved_case + S[i]) * (P[i] / 100) \
        + unsolved_case * ((100 - P[i]) / 100))

print(dp[0][X])