PythonでJOI難易度6を埋める #22

B - 古本屋 (Books)

DPDPDPDPDPDPDPDPDPDP．
DPで解く．ジャンルごとに売る本を決めるとき，同じジャンルの中で高いものから売った方が合計買取額は大きくなる．そこで，ジャンルごとに降順でsortし，買取ボーナスも含めた値段で累積和を取る．ジャンルの累積和をとするとき，の番目は，]]]] + となる．
ここまできたら，後はDPで解ける．]]=番目のジャンルまでで，合計冊になるように本を売ったときの最大の買取金額とする．このDPはナップサックDPの重さが冊数になったものなので，同じ要領で解ける．

from sys import stdin
input = stdin.readline
inf = float('inf')

n, k = map(int, input().strip().split())
books = [list(map(int, input().strip().split())) for _ in range(n)]

genre = [[] for _ in range(10)]
for i in range(n):
    genre[books[i][1] -1].append(books[i][0])

cnt_genre = 0
cnt_books = 0
for i in range(10):
    if len(genre[i]) != 0:
        cnt_genre += 1
        genre[i].sort(reverse = True)
        cnt_books = max(len(genre[i]), cnt_books)

csum_genre = [[] for _ in range(cnt_genre)]
for i in range(10):
    if len(genre[i]) == 0:
        continue
    for j in range(len(genre[i])):
        if j == 0:
            csum_genre[i].append(genre[i][j])
        else:
            csum_genre[i].append(genre[i][j] + csum_genre[i][-1] + 2 * j)

dp = [[0 for _ in range(max(cnt_books, k) + 1)] for _ in range(cnt_genre)]

for i in range(1, len(csum_genre[0]) + 1):
    dp[0][i] = csum_genre[0][i - 1]

for i in range(1, cnt_genre):
    length = len(csum_genre[i])
    for j in range(max(cnt_books, k) + 1): #ジャンルiまででj冊売る
        tmp = -inf
        if j == 0:
            dp[i][j] = 0
            continue
        if length - j >= 0:
            for l in range(j): #ジャンルiをl冊まで売る(0 <= l <= j) l+1冊売却
                tmp = max(dp[i - 1][j - l - 1] + csum_genre[i][l], tmp)
        else:
            for l in range(length):
                tmp = max(dp[i - 1][j - l - 1] + csum_genre[i][l], tmp)
        dp[i][j] = max(dp[i - 1][j], tmp, dp[i][j])
print(dp[-1][k])