树形DP Python 洛谷P1352

树形DP Python
洛谷P1352：
某大学有 n个职员，编号为 1~N。他们之间有从属关系，也就是说他们的关系就像一棵以校长为根的树，父结点就是子结点的直接上司。现在有个周年庆宴会，宴会每邀请来一个职员都会增加一定的快乐指数a[i] ，但是呢，如果某个职员的上司来参加舞会了，那么这个职员就无论如何也不肯来参加舞会了。所以，请你编程计算，邀请哪些职员可以使快乐指数最大，求最大的快乐指数。
输入：

7
1
1
1
1
1
1
1
1 3
2 3
6 4
7 4
4 5
3 5

输出：

5

代码：

n=int(input())    #公司人数
h=[]    #存储每个人的气氛值
a=[0]*(n)    #a[i]为i的上司
b={}    #b{上司：[下属1，下属2]}
dp=[[0]*2 for _ in range(n)]    #dp[i][0]表示i不去，dp[i][1]表示i去
flag=[]
for i in range(n):
    h.append(int(input()))
    flag.append(1) #指示没有boss的根位置
    b[i]=[] #初始化字典
for i in range(n-1):
    n1,n2=map(int,input().split())
    a[n1-1]=n2-1#所有职员编号=原编号-1，即从0开始计数
    b[n2-1].append(n1-1)
    flag[n1-1]=0#有上司的职员被赋值为0，最后非0的是根节点

def  dfs(x):
    global dp
    global b
    for j in range(len(b[x])):
        
        i=b[x][j]  #遍历x的孩子
        dfs(i)
        dp[x][0]+=max(dp[i][0],dp[i][1])
        dp[x][1]+=dp[i][0]
    dp[x][1]+=h[x] #快乐指数

for i in range(n):
    if (flag[i]==1):  #在根节点处开始
        
        dfs(i)
        print(max(dp[i][1],dp[i][0]))
        break

参考学习：
https://blog.youkuaiyun.com/liuzich/article/details/107672918
https://oi-wiki.org/dp/tree/