AcWing 树形DP相关问题 1074. 二叉苹果树

import sys
sys.stdin = open('data.txt', 'r')



N, Q = map(int, input().split())

link = {}
for i in range(1, N):
    a, b, w = map(int, input().split())
    if a not in link:
        link[a] = []
    if b not in link:
        link[b] = []

    link[a].append((b, w))
    link[b].append((a, w))


'''
在树上选择Q条边，等价于选择连在一起的Q+1个节点，每个节点的权值就是和其父节点连接的边的权值，
根节点的权值为0，问题就变成了一个有依赖的背包问题，首先根节点必选，如果一个节点被选择了，其
父节点必定被选择

dp(i, j) 表示以i为根的子树上选择j个节点，所有节点的权值最大和是多少
在进行决策时候其实就是两个分组进行背包，物品开销是子树选择节点数jj，物品的价值就是dp(子节点id, jj)

'''

# 每个节点的节点数和权重
node_num = [0] * (N+1)
weight = [0] * (N+1)


# 先dfs一把，把node_num, weight填充好
def dfs(root, prev = None):
    node_cnt = 1
    for child, w in link[root]:
        if child == prev:
            continue

        weight[child] = w
        node_cnt += dfs(child, root)

    link[root] = [val for val in link[root] if val[0] != prev]
    node_num[root] = node_cnt
    return node_cnt

dfs(1)


from functools import lru_cache
@lru_cache(typed=False, maxsize=128000000)
def dp(i, j):
    if j == 0:
        return 0

    if j > node_num[i]:
        # 不可能存在方案
        return -0x7fffffff

    if j == 1:
        return weight[i]

    max_val = 0
    left_id, right_id = link[i][0][0], link[i][1][0]

    for left_j in range(0, j):
        right_j = j-1-left_j
        max_val = max(max_val, dp(left_id, left_j) + dp(right_id, right_j))
    return max_val + weight[i]

print(dp(1, Q+1))