小扣有一个根结点为 root 的二叉树模型,初始所有结点均为白色,可以用蓝色染料给模型结点染色,模型的每个结点有一个 val 价值。小扣出于美观考虑,希望最后二叉树上每个蓝色相连部分的结点个数不能超过 k 个,求所有染成蓝色的结点价值总和最大是多少?
示例 1:
输入:root = [5,2,3,4], k = 2
输出:12
解释:结点 5、3、4 染成蓝色,获得最大的价值 5+3+4=12
示例 2:
输入:root = [4,1,3,9,null,null,2], k = 2
输出:16
解释:结点 4、3、9 染成蓝色,获得最大的价值 4+3+9=16
方法:动态规划
题意分析
找到染色节点的最大和
思路及算法
由于有下一步的计算结果依赖于前一步,可尝试用动态规划的方式求解。
一般的动态规划题目思路三步走:
定义状态转移方程
给定转移方程初始值
写代码递推实现转移方程
考虑自底向上dp,保存子节点的状态并计算当前是否染色的状态。
1. 定义状态转移方程
定义一维数组 dp[k+1]:
dp[i] 表示 每个节点的状态,i 表示染了几个节点,i=0 表示没有染色,i>0 表示染色 。
定义状态转移方程:
当前节点为root,dp逻辑为(详见注释):
root不染色,那么只要返回 dp[0],其值为左、右子树染色或不染色的最大值之和
root染色,那么就分左子树染色 j 个,右子树染色 i - 1 - j 个时,加上 root.val 的和。
注意:j 需要从 0 取到 i - 1,也就是包含 l[0] 和 r[0]。因为 l[0] 也包含左子树染了j个节点的情况,因为左子树的下一层子节点可能染了j个节点。
dp[i] = Math.max(dp[i], root.val + l[j] + r[i - 1 - j]);
条件分叉代码可参考注释。
2. 给定转移方程初始值
初始值为 叶子节点下的空节点,所以 直接返回都为 0 的 new int[k + 1];
class Solution { public int maxValue(TreeNode root, int k) { int[] dp = dynamic(root, k); int max = Integer.MIN_VALUE; for (int i = 0; i <= k; i++) { //取root的各种染色情况的最大值 max = Math.max(max, dp[i]); } return max; } private int[] dynamic(TreeNode root, int k) { int[] dp = new int[k + 1]; //1.初始化:空节点为底,自底向上 if (root == null) return dp; //2.获取左、右子树染色状态的dp表 //- 左子树 int[] l = dynamic(root.left, k); //- 右子树 int[] r = dynamic(root.right, k); //3.更新处理root 染色/不染色 的情况下的dp表 //- 不染root int ml = Integer.MIN_VALUE, mr = Integer.MIN_VALUE; for (int i = 0; i <= k; i++) { //- 分别取子节点的最大值 ml = Math.max(ml, l[i]); mr = Math.max(mr, r[i]); } dp[0] = ml + mr; //- 染root for (int i = 1; i <= k; i++) { for (int j = 0; j < i; j++) { //- 还需要染色 i - 1 个点,左子树 j 个,右子树 i-1-j 个 dp[i] = Math.max(dp[i], root.val + l[j] + r[i - 1 - j]); } } //4.更新完毕,返回后继续向上动态规划 return dp; } }