abc 060 D-Simple Knapsack (python)

タイトルからして dp だと思ったら解説は全探索だった。本質的には dp もメモ化全探索みたいなところあるけど。
w_i の制約が特徴的で max(w_i) - w_1 <= 3 とある。なので
dp[ i ][ j ][ k ] で i 番目の荷物まで見たときに、j 個の荷物を入れて、
重さが j * w_1 + k となっているときの価値の最大値
として dp を作ってみた。
（k は選んだ荷物の重さの合計 - 選んだ数 j * 荷物の最小値 w_1で与えられる）
そうすることで重さが w を超えるような選び方を排除しつつ、
最大 100*100*(3*99) で計算できるようになった。
（荷物の数は100以下で、k が最大となるのは w_2,...,w_n が全て w_1+3 となる時）
最大値は dp[ n ] のどこかにあるので、思考停止して dp[ n ] を全て展開して最大値を求めた。
初期値を -float('inf') にしているが、dp[ i ][ 0 ][ 0 ] だけは全て 0 にできるように実装するのに注意。

import sys
input = sys.stdin.readline
import itertools
n,w = map(int,input().split())
c = []
s = 0
for i in range(n):
  a,v = map(int,input().split())
  s += a
  c.append((a,v))
  
r = c[0][0]
L = min(100,w//r)
m = s-r*n+1
dp = [[[-float('inf')]*m for i in range(L+1)] for i in range(n+1)]
for i in range(n+1):
  dp[i][0][0] = 0

for i in range(n):
  for j in range(L):
    for k in range(m):
      W,V = c[i]
      R = W-r
      if k-R >= 0 and r*(j+1)+k <= w:
        dp[i+1][j+1][k] = max(dp[i][j+1][k],dp[i][j][k-R]+V)
      else:
        dp[i+1][j+1][k] = dp[i][j+1][k]
        
print(max(list(itertools.chain.from_iterable(dp[n]))))