AcWing 树形DP相关问题 1077. 皇宫看守

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




from functools import lru_cache
N  = int(input())

link = {}
weight = {}
parent = {}
for _ in range(N):
    arr = list( map(int, input().split()) )
    link[arr[0]] = arr[3:]
    weight[arr[0]] = arr[1]
    for sub in arr[3:]:
        parent[sub] = arr[0]

# 以i为根的子树上进行选择，且i节点的状态是j的情况下，所哟选择中最小的开销
# 节点的状态可以分成3中状态
# 0. 自己不放警卫，但是被父节点看到
# 1. 自己不放警卫，但是被子节点看到
# 2. 自己放警卫，被自己看到
# 前面两种状态集合有重叠，但是不影响求解
@lru_cache(typed=False, maxsize=128000000)
def dp(i, j):
    if j == 0:
        # 自己已经被父节点看到，那子节点 可以放警卫，也可以不放警卫但是被子节点看到
        ans = 0
        for child in link[i]:
            ans += min(dp(child, 2), dp(child, 1))
        return ans

    elif j == 1:
        # 自己不放警卫，但是被子节点看到，那所有子节点至少要有一个放警卫
        if len(link[i]) == 0:
            return 0x7fffffff   # 没有可行方案

        # 下面开始背包，每个子节点看成一个分组，每个分组有两个物品，分别是放警卫和不放警卫，要求至少有一个分组要选放警卫
        n = len(link[i])
        children = link[i]
        f = [[0, 0] for _ in range(n)]      # f(ii, jj)表示前ii个子节点至少jj个选择放警卫的所有选法中的最小总开销

        for ii in range(n):
            for jj in range(2):
                if ii == 0:
                    if jj == 1:
                        f[ii][jj] = dp(children[ii], 2)
                    else:
                        f[ii][jj] = min(dp(children[ii], 2), dp(children[ii], 1))
                else:
                    f[ii][jj] = min( f[ii-1][jj] + dp(children[ii], 1), f[ii-1][jj-1] + dp(children[ii], 2) )

        return f[n-1][1]

    else:
        # 自己放警卫，那子节点选择就是要么自己也放一个警卫，要么不放警卫但是被父节点看到
        ans = weight[i]
        for child in link[i]:
            ans += min(dp(child, 2), dp(child, 0))

        return ans


# 找到根节点
root = 1
while root in parent:
    root = parent[root]

print( min(dp(root, 2), dp(root, 1)) )