🐸

[鉄則A18] Subset Sum

2024/03/25に公開

https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_r

二次元の動的計画法の基本が学べる部分和問題。

問題の要約

  • N 枚のカードから何枚か選択したとき合計が S になるか?

例1で問題を理解する

2 2 3 の値が書かれた 3 枚のカードから何枚か選択したとき合計 7 になるか? 答えは Yes で、すべてを選択したとき合計 7 になる。

全探索な解き方 (TLE)

N, S = gets.split.collect(&:to_i)  # => [3, 7]
A = gets.split.collect(&:to_i)     # => [2, 2, 3]
success = (1..A.size).any? do |max|
  A.combination(max).any? { |cards| cards.sum == S }
end
success                            # => true
puts success ? "Yes" : "No"

合計値(S)は 1 以上と決まっているので最低一枚は選択する必要がある。したがって 1 から 3 枚を選択するすべてパターンを順に試す。この方法は N が小さいときにだけ使える。

二次元な動的計画法 (AC)

N, S = gets.split.collect(&:to_i)  # => [3, 7]
A = gets.split.collect(&:to_i)     # => [2, 2, 3]
A.unshift(nil)                     # => [nil, 2, 2, 3]
dp = Array.new(N.next) { Array.new(S.next) }
dp[0][0] = true
(1..N).each do |y|                  # 縦軸: 1..カード枚数
  (0..S).each do |x|                # 横軸: 0..合計値
    if x < A[y]
      # y 番目のカードが使えない
      dp[y][x] = dp[y - 1][x]
    else
      # y 番目のカードを使えるが、使わない or 使う
      dp[y][x] = dp[y - 1][x] || dp[y - 1][x - A[y]]
    end
  end
end
success = dp[N][S]                 # => true
puts success ? "Yes" : "No"
0 1 2 3 4 5 6 7
0 (カード値)
1 (2)
2 (2)
3 (3)
  • x < A[y] なら現在のカードを選択できない (カード値に足りていないため)
  • x >= A[y] なら現在のカードを選択している可能性がある
  • 上からコピーするのは現在のカードを使わなかったとする場合
  • 左上からコピーするのは現在のカードを使ったとする場合
  • 現在のカードを使った想定のとき、元は現在のカードを使わなかった地点 x - A[y] になる
  • 「使わなかった場合」と「使った場合」の順番は関係なくどっちかが真なら印をつける
  • dp[2][2] の○からは 1, 2枚目のどちらを使ったのか判断できない
  • dp[3][7] の○は左上の dp[2][4] を経由しているので2枚目のカードを使っている
  • 左上から経由したカードだけが使用されている
  • dp[3][7] は左上方向に dp[2][4]dp[1][2] と戻れるので 3, 2, 1 枚目を使っている

緩和部分のリファクタリング

鉄則本の解説では次のような感じで書かれていて冗長だがわかりやすい。

if x < A[y]
  # y 番目のカードが使えない
  dp[y][x] = dp[y - 1][x]
end
if x >= A[y]
  # y 番目のカードを使えるが、使わない or 使う
  dp[y][x] = dp[y - 1][x] || dp[y - 1][x - A[y]]
end

ここを DRY にするなら、

dp[y][x] = dp[y - 1][x] || (x >= A[y]) && dp[y - 1][x - A[y]]

と書ける。

分岐条件まとめ

条件 コード 向き
使えない x < A[y] && dp[y][x] = dp[y - 1][x] ⬇️
使えるが使わない x >= A[y] && dp[y][x] = dp[y - 1][x] ⬇️
使えるので使う x >= A[y] && dp[y][x] = dp[y - 1][x - A[y]] ↘️

Discussion