Python动态规划——以“加分二叉树”为例

题目链接：登录—专业IT笔试面试备考平台_牛客网

题目描述

设一个n个节点的二叉树tree的中序遍历为（l,2,3,…,n），其中数字1,2,3,…,n为节点编号。每个节点都有一个分数（均为正整数），记第j个节点的分数为di，tree及它的每个子树都有一个加分，任一棵子树subtree（也包含tree本身）的加分计算方法如下：

    subtree的左子树的加分× subtree的右子树的加分＋subtree的根的分数

    若某个子树为主，规定其加分为1，叶子的加分就是叶节点本身的分数。不考虑它的空子树。 试求一棵符合中序遍历为（1,2,3,…,n）且加分最高的二叉树tree。

要求输出：

    （1）tree的最高加分

    （2）tree的前序遍历

输入描述 

第1行：一个整数n（n＜30），为节点个数。
第2行：n个用空格隔开的整数，为每个节点的分数（分数＜100）。 

输出描述 

第1行：一个整数，为最高加分（结果不会超过4,000,000,000）。
第2行：n个用空格隔开的整数，为该树的前序遍历。 

 示例

输入 

5
5 7 1 2 10

输出 

145
3 1 2 4 5

思路 

       首先需要知道什么是二叉树中的中序遍历，即一颗二叉树按照左父右的顺序输出节点，题目给出中序遍历为1，2，3，…，n，要根据这个排序进行构造二叉树从而得到最大的分数。根据这个排序，每个节点都可以作为最初的根节点，例如1，2，3，4，5，6，7这个中序遍历，将4作为根节点，则(1,2,3)为4的左子树，（5，6，7)为4的右子树，以此类推，直到用完所有节点。因此遍历所有节点为根节点，再根据加分计算方法：subtree的左子树的加分× subtree的右子树的加分＋subtree的根的分数，可以得到以下状态转移方程。

    其中dp[i][j]表示从节点i到节点j的最大分数，节点k表示当前根节点，dp[i][k - 1]为左子树中节点i到节点k - 1的最大分数，dp[i][k + 1]为右子树中节点i到节点k + 1的最大分数，最后加上根节点k的分数。

    题目还要求输出最大分数的二叉树的前序遍历，因此用root[i][j]记录从节点i到节点j的最大分数时的根节点，有了每一段的根节点就可以按照前序遍历输出。

代码 

第一段代码：对数据进行读取以及定义变量。考虑到在状态转移方程中k为1时即最前面的节点作为根节点，则没有左子树，根据题目要求不考虑空子树则设空子树的分数为1，不影响左子树的分数与右子树的分数进行相乘，因此设置长度宽度皆为n + 1的dp表防止因为边界问题而报错。对于dp[i][i]即节点i到节点i的初始分数则是本身，初始根节点也是本身，进行赋初值。

n = int(input())
a = [0] + list(map(int,input().split()))
dp = [[1] * (n + 1) for _ in range(n + 1)]
root = [[0] * (n + 1) for _ in range(n + 1)]
for i in range(1,n + 1):
    dp[i][i] = a[i]
    root[i][i] = i

 第二段代码：动态规划的核心部分。这里我们按区间长度2到n进行遍历，其中再遍历区间包含的所有节点对dp表进行更新。由于刚开始只知道每个根节点(1,1),(2,2),(3,3)的分数，例如在求dp[1][2]时需要先求得dp[1][1]和dp[2][2]，求dp[1][3]时需要先求得dp[2][3]，因此按照阶梯式即区间长度对dp表进行更新，可以避免用未知数据计算得出的结果。其中在对dp表进行更新时，每次从新的一层进行更新时，最初的节点会作为第一个根节点，因此需要对这个根节点赋初值，即节点i到节点j的最大分数 = 节点i + 1到节点j的最大分数 + 根节点i的分数，即考虑遍历根节点初始左子树为空的情况。

for length in range(2,n + 1):
    for i in range(1,n - length + 2):
        j = i + length - 1
        dp[i][j] = dp[i + 1][j] + dp[i][i]
        root[i][j] = i
        for k in range(i,j):
            if dp[i][j] < dp[i][k - 1] * dp[k + 1][j] + dp[k][k]:
                dp[i][j] = dp[i][k - 1] * dp[k + 1][j] + dp[k][k]
                root[i][j] = k

print(dp[1][n])
tree(1,n)

 第三段代码：一个输出树的前序遍历的函数。终止条件则是左节点大于右节点时则代表当前的树已经结束。

def tree(l,r):
    if l > r:
        return
    print(root[l][r],end = ' ')
    tree(l,root[l][r] - 1)
    tree(root[l][r] + 1,r)

完整代码如下。

def tree(l,r):
    if l > r:
        return
    print(root[l][r],end = ' ')
    tree(l,root[l][r] - 1)
    tree(root[l][r] + 1,r)

n = int(input())
a = [0] + list(map(int,input().split()))
dp = [[1] * (n + 1) for _ in range(n + 1)]
root = [[0] * (n + 1) for _ in range(n + 1)]
for i in range(1,n + 1):
    dp[i][i] = a[i]
    root[i][i] = i

for length in range(2,n + 1):
    for i in range(1,n - length + 2):
        j = i + length - 1
        dp[i][j] = dp[i + 1][j] + dp[i][i]
        root[i][j] = i
        for k in range(i,j):
            if dp[i][j] < dp[i][k - 1] * dp[k + 1][j] + dp[k][k]:
                dp[i][j] = dp[i][k - 1] * dp[k + 1][j] + dp[k][k]
                root[i][j] = k

print(dp[1][n])
tree(1,n)

回顾 

    在做这题时也花费了很多的时间，问题在于对树的不熟练以及该如何更新dp表。起初的思路有DFS、BFS等等，我想是找最大值以及路径，用DFS应该是可以实现的，动态规划本质上也是个搜索，不断更新dp表的一个过程。代码是比较短的，但其中的思路并不好想，包括如何定义状态转移方程、如何更新dp表对于目前的我来说还是存在一定的困难的。经常在比赛中会看到树的身影，包括刚开始在做这题时会想到用队列来存储树的每一层的一个根节点和最大分数，之后还是要加强对基本数据结构的一个理解和运用。