AcWing 背包相关问题 487. 金明的预算方案

'''
分组DP

按照主件进行分组，每一组里面选物品组合，只有五种情况进行决策
1 一个都不选
2 只选主件
3 选主件和附件1
4 选主件和附件2
5 选主件和附件1和附件2
'''

N, m = map(int, input().split())
g = {}
for i in range(1, m + 1):
    v, p, q = map(int, input().split())
    grp_id = q if q != 0 else i
    if grp_id not in g:
        g[grp_id] = [None, []]

    if q == 0:
        g[grp_id][0] = (v, p)
    else:
        g[grp_id][1].append((v, p))

# arr中保存每一个分组的信息
arr = [val for val in g.values()]
# print(arr)


# dp(i, j) 表示前i个分组做选择，总费用不超过j约束下所有选择中总价值的最大值
dp = [0] * (N + 1)

for i in range(len(arr)):
    cost_and_value = [(arr[i][0][0], arr[i][0][0] * arr[i][0][1])]
    for cost, imp in arr[i][1]:
        cost_and_value.append((arr[i][0][0] + cost, arr[i][0][0] * arr[i][0][1] + cost * imp))
    if len(arr[i][1]) == 2:
        cost1, imp1 = arr[i][1][0]
        cost2, imp2 = arr[i][1][1]
        cost_and_value.append((arr[i][0][0] + cost1 + cost2, arr[i][0][0] * arr[i][0][1] + cost1 * imp1 + cost2 * imp2))

    for j in range(N, -1, -1):
        if i == 0:
            for cost, value in cost_and_value:
                if j >= cost:
                    dp[j] = max(dp[j], value)
        else:
            for cost, value in cost_and_value:
                if j >= cost:
                    dp[j] = max(dp[j], dp[j - cost] + value)

print(dp[N])