[Python] Educational DP Contest E – Knapsack 2

問題

E – Knapsack 2

回答

TLEで間に合わない最初の回答

アルゴリズム的には合っているかな?

import sys

# input処理を高速化する
input = sys.stdin.readline

# 入力
N, W = map(int, input().split())
# sum_vの取りうる上限値
MAX_V = (10**3+1) * N

# 最小値を求めるので無限大の値
f_inf = float('inf')

def chmin(a, b):
    if a <= b:
        return a
    else:
        return b

def knapsack(n, w):
    # dpの初期化 dp[i][sum_v]は重さの総和の最小値
    dp = [[f_inf for i in range(MAX_V)] for j in range(N+1)]
    # 初期条件
    dp[0][0] = 0
    # forは全て関数の中で回す。
    # 入力
    for i in range(N):
        w, v = map(int, input().split())
        for sum_v in range(MAX_V):
            # i番目の品物を選ぶ場合
            if (sum_v -v >= 0):
                dp[i + 1][sum_v] = chmin(dp[i + 1][sum_v], dp[i][sum_v - v] + w)
            # i番目の品物を選ばない場合
            dp[i + 1][sum_v] = chmin(dp[i + 1][sum_v], dp[i][sum_v])
    #dp[N][sum_v]の値が、W以下であるsum_vの値の最大値
    ans = 0
    for sum_v in range(MAX_V):
        if dp[N][sum_v] <= W:
            ans = sum_v
    print(ans)

knapsack(N, W)

AC

少し絞るとACになる。

import sys

# input処理を高速化する
input = sys.stdin.readline

# 入力
N, W = map(int, input().split())
# sum_vの取りうる上限値
MAX_V = (10**3) * N +1
# 最小値を求めるので無限大の値
f_inf = float('inf')

def knapsack(n, w):
    # dpの初期化 dp[sum_v]は重さの総和の最小値
    dp = [f_inf] * MAX_V
    # 初期条件
    dp[0] = 0
    sum_v = 0
    # forは全て関数の中で回す。
    # 入力
    for i in range(N):
        w, v = map(int, input().split())
        for j in range(MAX_V, v-1, -1):
            new_weight = dp[j - v] + w
            if new_weight > W:
                continue
            if dp[j] > new_weight:
                dp[j] = new_weight

    ans = 0
    for sum_v in range(MAX_V):
        if dp[sum_v] <= W:
            ans = sum_v
    print(ans)

knapsack(N, W)

numpyを使った回答

import sys
import numpy as np

# input処理を高速化する
input = sys.stdin.readline

# 入力
N, W = map(int, input().split())
# sum_vの取りうる上限値
MAX_V = (10**3) * N +1

def knapsack(n, w):
    # dp[i][j] : i番目の品物で価値jの重さの最小値
    dp = np.full([N+1, MAX_V], np.inf)
    dp[0][0] = 0
    for i in range(n):
        w, v = map(int, input().split())
        dp[i+1] = dp[i]
        dp[i+1][v:] = np.minimum(dp[i][v:], dp[i][:-v] + w)
    print(np.max(np.where(dp <= W)))

knapsack(N, W)