蓝桥 ALGO-4 结点选择 python

问题描述

有一棵 n 个节点的树，树上每个节点都有一个正整数权值。如果一个点被选择了，那么在树上和它相邻的点都不能被选择。求选出的点的权值和最大是多少？

输入格式

第一行包含一个整数 n 。

接下来的一行包含 n 个正整数，第 i 个正整数代表点 i 的权值。

接下来一共 n-1 行，每行描述树上的一条边。

输出格式

输出一个整数，代表选出的点的权值和的最大值。

样例输入

5
1 2 3 4 5
1 2
1 3
2 4
2 5

样例输出

12

样例说明

选择3、4、5号点，权值和为 3+4+5 = 12 。

数据规模与约定

对于20%的数据， n <= 20。

对于50%的数据， n <= 1000。

对于100%的数据， n <= 100000。

权值均为不超过1000的正整数。

结果：70%，最后三个100000数量级的数据报错，其他正确，估计是数组过大导致。不知道从何改进。

思路：使用树形动态规划，也就是在树的结构上进行循环递归，更新动态规划的m值。

m[i][1]代表在以i为根节点的子树中选择i节点所得最优结果。m[i][0]代表不选择i节点的最优结果

dp(a,b)用来更新a节点的m值，也就是m[a]，其中b为a的父节点

当节点i为叶子节点时，m[i][0] = 0, m[i][1] = value[i]

当节点i为非叶子节点时，m[i][0] = Σ ( max ( m[i的子节点][1] , m[i的子节点][0] ) )

                                         m[i][1] = Σ  m[i的子节点][0]

n = int(input())#n为顶点个数
value = list(map(int,input().split()))#每个节点价值
tree = []
#用来记录树，二维数组，举例：[[],[2,4],[1,5,8],[9]],tree[1]得到的是与第一个节点连接的其他点。
for i in range(n):#开辟空间
    tree.append([])
for i in range(n-1):
#初始化tree，按照下面的保存方法，则tree[i]得到的是第i个节点的所有孩子节点和一个父节点
    a,b = map(int,input().split())
    tree[a].append(b)
    tree[b].append(a)
#已知输入必是一棵树，所以可以以任意一个节点为根节点，这里选择第一个为根节点
m = []
#m[i][1]表示的是在以i为根节点的子树中，选择i节点的最优结果，m[i][0]是不选择i节点的最优结果。
for i in range(n+1):
#初始化m，使所有节点i的m[i][1]都初始为value[i]，这是因为m[i][1]表示的是在以i为根节点的子树中选
#择i节点的最优结果，那么m[i][1]=value[i]+所有孩子节点的为根的子树的m[i][0],初始为value[i]之
#后，在递归动态规划中进行累加。
#m[i][0]初始化为0，如果第i个节点不选，其子节点可以选也可以不选，所以m[i][0]=所有孩子节点的为根
#的子树的max(m[i][0],m[i][1])的累加
    m.append([])
    m[i].append(0)
    m[i].append(value[i-1])

def dp(a,b):#a代表的是此节点的m[i][0],m[i][1]要更新，b代表是a节点的父节点
    if len(tree[a])==1:#代表此节点为叶子节点，则m[i][1],m[i][0]与初始值相同，不必改动
        return
    else:
        for i in range(len(tree[a])):
            chile = tree[a][i]#循环与a节点连接的所有点
            if chile != b:
#如果不是父节点。若要m[i]则必须知道子节点的m[i],但是chile包含了所有连接的点，所以要跳过父节点
                dp(chile,a)#先更新孩子节点为根节点的子树的m值
                m[a][1] += m[chile][0]#累加
                m[a][0] += max(m[chile][1],m[chile][0])

dp(1,0)
#任意选择节点为整颗树的根节点，这里选择1，表示第一个节点，因为节点是从1开始，0位置不代表节点
print(max(m[1][0],m[1][1]))打印更新后的结果