LeetCode——二叉树染色

题目描述

小扣有一个根结点为 root 的二叉树模型，初始所有结点均为白色，可以用蓝色染料给模型结点染色，模型的每个结点有一个 val 价值。小扣出于美观考虑，希望最后二叉树上每个蓝色相连部分的结点个数不能超过 k 个，求所有染成蓝色的结点价值总和最大是多少？
示例 1：
输入：root = [5,2,3,4], k = 2
输出：1

解释：结点 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

提示：
1 <= k <= 10
1 <= val <= 10000
1 <= 结点数量 <= 10000

题解

• 根据K判断当前节点可不可以染。

• 可以染的情况下还需判断染或不染。

• 使用Map记录每个节点不染情况下，以其为根节点的子树的最大收益。避免超时

  /**
  * Definition for a binary tree node.
  * public class TreeNode {
  *     int val;
  *     TreeNode left;
  *     TreeNode right;
  *     TreeNode(int x) { val = x; }
  * }
  */
  class Solution {
    public int maxValue(TreeNode root, int k) {
    // map用于记录当前根节点不染时，子树的最大收益
    HashMap<TreeNode,Integer>map=new HashMap<>();
    return solve(root,k,k,map);
    }
    public int solve(TreeNode root,int k,int curK,HashMap<TreeNode,Integer>map) {
    if(root==null) {
        return 0;
    }
    // 连续染色的节点数达到阈值，当前节点不染
    if(curK==0) {
        if(map.containsKey(root)) {
            return map.get(root);
        }
        int value=solve(root.left,k,k,map)+solve(root.right,k,k,map);
        map.put(root, value);
        return value;
    }
    int value1=0;
    // 当前节点染时，左右连续被染的节点数不能超过阈值
    for(int left=0;left<=curK-1;++left) {
        int value=root.val+solve(root.left,k,left,map)+solve(root.right,k,curK-1-left,map);
        if(value1<value) {
            value1=value;
        }
    }
  
    // 连续被染节点数没达到阈值时，当前节点可以不染
    int value2=0;
    if(map.containsKey(root)) {
        value2= map.get(root);
    }else {
        int valueLeft=solve(root.left,k,k,map);
        int valueRight=solve(root.right,k,k,map);
        value2=valueLeft+valueRight;
        map.put(root, value2);
    }
    return value1>value2?value1:value2;
    }
  }