LCP 34. 二叉树染色 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Solution {
    public int maxValue(TreeNode root, int k) {
        // 定义 dp[i] 为以当前节点为根，所能到达的所有相邻的蓝色节点个数（包含当前节点）为 i 的情况下，节点价值总和的最大值
        int[] dp = getDp(root, k);

        int ans = Integer.MIN_VALUE;
        // 注意此处的 for 循环，可以使用树 [8,1,3,9,9,9,null,9,5,6,8] 加深理解
        for (int i = 0; i < dp.length; i++) {
            ans = Math.max(ans, dp[i]);
        }
        return ans;
    }

    private int[] getDp(TreeNode root, int k) {
        int[] dp = new int[k + 1];
        if (root == null) {
            return dp;
        }

        int[] leftDp = getDp(root.left, k);
        int[] rightDp = getDp(root.right, k);

        // 当不将 root 染成蓝色时，可推出 dp[0] 的值，此时 dp[0] 最大为左侧子树和右侧子树各自最大值的和
        int leftMax = Integer.MIN_VALUE, rightMax = Integer.MIN_VALUE;
        for (int i = 0; i <= k; i++) {
            leftMax = Math.max(leftMax, leftDp[i]);
            rightMax = Math.max(rightMax, rightDp[i]);
        }
        dp[0] = leftMax + rightMax;

        // 当将 root 染成蓝色时，此时左子树和右子树最多只能连续染 i - 1 个节点
        for (int i = 1; i < dp.length; i++) {
            for (int leftCount = 0; leftCount < i; leftCount++) {
                int rightCount = i - leftCount - 1;
                dp[i] = Math.max(dp[i], leftDp[leftCount] + root.val + rightDp[rightCount]);
            }
        }

        return dp;
    }
}

References

LCP 34. 二叉树染色