[Python] 動的計画法 ABC015D

#ABC015D

個数制限付きのナップサック問題です。ちなみに全探索で考えると、N 個の中から K 個選んで入れる・入れないを全て試すと $O(2^N)$ となり、間に合いません。動的計画法であれば、計算量は $O(NKW)$ で間に合います。

動的計画法は部分問題を意識して、次のように設計します。

$dp[n][k][w]$の定義：
n 番目以下のスクリーンショットから k 枚以下のものを選び、幅が w を超えないように荷物を選んだ時の重要度の最大値

dpの漸化式：
$\mathrm {\rm dp}(n, k, w)$

=
\left \{
    \begin{array}{l}
        \mathrm {\rm max}\{ {\rm dp}( n - 1, k, w ), \:{\rm dp}( n - 1, k - 1, w - A_{n-1} ) + B_{k-1} \}&( k \geq A_{n-1} )\\
        \mathrm {\rm dp}( n - 1, k, w )　&(otherwise)\\
    \end{array}
\right.

dp初期条件：
$dp[0][0][*]=0$

求めるもの：
$dp[N][K][W]$

実装は、ボトムアップ方式の貰うDPで行います。Python では TLE ですが、PyPy では間に合いました。

サンプルコード

W = int(input())
N, K = map(int, input().split())
AB = [list(map(int, input().split())) for _ in range(N)]
dp = [[[0]*(W+1) for _ in range(K+1)] for _ in range(N+1)]
for n in range(N):
    for k in range(K+1):
        for w in range(W+1):
            a = AB[n][0] # 幅
            b = AB[n][1] # 重要度
            if w-a>=0 and k-1>=0:
                dp[n+1][k][w] = max(dp[n][k][w], dp[n][k-1][w-a]+b)
            else:
                dp[n+1][k][w] = dp[n][k][w]
print(dp[N][K][W])