アルゴリズム実技検定(PAST) 第11回 H問題 Python解答例(ナップサック問題 / 動的計画法)

Supershipの名畑です。
うる星やつらの2クール目のOPも好きです。テンちゃんと竜之介が加わると一気にうる星やつら感が増します。

はじめに

アルゴリズム実技検定(PAST)の過去問シリーズです。
第11回H問題の解答例です。

難易度としてはABC(AtCoder Beginner Contest)のC〜Dぐらいという肌感ですが、ナップサック問題 / 動的計画法を知っていればほぼノータイムで解けそうではあります。

「そもそもアルゴリズム実技検定ってなんですか？」という方はこちらの記事を合わせてご覧いただけますと幸いです。

第11回 H問題 2つのナップサック

問題

H - 2つのナップサック

言語

• PyPy3 (7.3.0)

Python(3.8.2)だとTLEがとれなかったのでPyPyを使いました。

解答例

有名なナップサック問題です。
注意点としては、ナップサックが2つあるため動的計画法(Dynamic Programming 略してDP)に用いる配列を三次元配列にしている点です。

N, A, B = map(int, input().split())

# 入力値wとvを配列に格納
w = [[] for i in range(N)]
v = [[] for i in range(N)]
for i in range(0, N):
    w[i], v[i] = map(int, input().split())

# 動的計画法のために三次元配列dp[N + 1][A + 1][B + 1]を用意
# 値（美味しさ）の初期値は0とする。
dp = [[[0 for k in range(B + 1)] for j in range(A + 1)] for i in range(N + 1)]

for i in range(N):
    for w1 in range(A + 1):
        for w2 in range(B + 1):
            # i番目のお菓子をいずれかのナップサックに入れられる
            if w[i] <= w1 or w[i] <= w2:
                # 1つ目のナップサックに入れられる
                if w[i] <= w1:
                    dp[i + 1][w1][w2] = max(dp[i][w1 - w[i]][w2] + v[i], dp[i][w1][w2])

                # 2つ目のナップサックに入れられる
                # 1つ目のナップサックに入れられたケースを考慮してdp[i + 1][w1][w2]もmaxの引数に追加
                if w[i] <= w2:
                    dp[i + 1][w1][w2] = max(dp[i][w1][w2 - w[i]] + v[i], dp[i][w1][w2], dp[i + 1][w1][w2])

            # i番目のお菓子を入れられないためi-1番目格納時点の美味しさのままとする
            else:
                dp[i + 1][w1][w2] = dp[i][w1][w2]

print(dp[N][A][B])

動的計画法とは、段階を踏んだ場合分けをして計算結果を記録して解いていく手法です。
今回の場合ですと

• 1つ目のループ(i)：お菓子が0番目、1番目、2番目……N-1番目(0からなので合計N個です)
• 2つ目のループ(w1)：1つ目のナップサックに入れられる重さの上限が0、1、2……A
• 3つ目のループ(w2)：2つ目のナップサックに入れられる重さの上限が0、1、2……B

という3重のループでの総当たりです。
そして最終的に dp[N][A][B] には「お菓子がN-1番目まで」「1つ目のナップサックに入れられる重さの上限がA」「2つ目のナップサックに入れられる重さの上限がB」という問題分の条件通りの結果が格納されることになります。

もしもナップサックが1つなら

ナップサックが1つの時はこんなコードです。
シンプルな方が考え方を理解しやすいかと思いますので参考までに。

# 入力
N, A = map(int, input().split())
w = [[] for i in range(N)]
v = [[] for i in range(N)]
for i in range(0, N):
    w[i], v[i] = map(int, input().split())

# dp用の二次元配列準備
dp = [[0 for i in range(A + 1)] for j in range(N + 1)]

for i in range(N):
    for w1 in range(A + 1):
        if w[i] <= w1:
            dp[i + 1][w1] = max(dp[i][w1 - w[i]] + v[i], dp[i][w1])
        else:
            dp[i + 1][w1] = dp[i][w1]

print(dp[N][A])

宣伝

SupershipのQiita Organizationを合わせてご覧いただけますと嬉しいです。他のメンバーの記事も多数あります。

Supershipではプロダクト開発やサービス開発に関わる方を絶賛募集しております。
興味がある方はSupership株式会社 採用サイトよりご確認ください。

是非ともよろしくお願いいたします。