LeetCode LCP 34. 二叉树染色（树上DP）

1. 题目

小扣有一个根结点为 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

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/er-cha-shu-ran-se-UGC
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

2. 解题

• 自底向上 DP
• unordered_map<TreeNode*, unordered_map<int, int>> m; 定义：每个节点TreeNode*，该节点相连的蓝色点数量，最大的和

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
    unordered_map<TreeNode*, unordered_map<int, int>> m;//每个节点TreeNode*，该节点相连的蓝色点数量，最大的和
    int n;
public:
    int maxValue(TreeNode* root, int k) {
        n = k;
        dfs(root);
        int ans = 0;
        for(auto it = m[root].begin(); it != m[root].end(); ++it)
        {
            int v1 = it->second;
            ans = max(ans, v1);
        }
        return ans;
    }
    void dfs(TreeNode* root)
    {
        if(!root)  return;
        dfs(root->left);
        dfs(root->right);
        if(m.count(root->left) && m.count(root->right))
        {
            for(auto it = m[root->left].begin(); it != m[root->left].end(); ++it)
            {
                int n1 = it->first, v1 = it->second;
                for(auto it1 = m[root->right].begin(); it1 != m[root->right].end(); ++it1)
                {
                    int n2 = it1->first, v2 = it1->second;
                    // root 不涂色，root相连的有色节点为0
                    m[root][0] = max(m[root][0], v1+v2);
                    if(n1+n2 < n)
                    {   // root 涂色
                        m[root][n1+n2+1] = max(m[root][n1+n2+1], v1+v2+root->val);
                    }
                }

            }
        }
        else if(m.count(root->left))
        {
            for(auto it = m[root->left].begin(); it != m[root->left].end(); ++it)
            {
                int n1 = it->first, v1 = it->second;
                // root 不涂色，root相连的有色节点为0
                m[root][0] = max(m[root][0], v1);
                if(n1 < n)
                {   // root 涂色
                    m[root][n1+1] = max(m[root][n1+1], v1+root->val);
                }

            }
        }
        else if(m.count(root->right))
        {
            for(auto it = m[root->right].begin(); it != m[root->right].end(); ++it)
            {
                int n1 = it->first, v1 = it->second;
                // root 不涂色，root相连的有色节点为0
                m[root][0] = max(m[root][0], v1);
                if(n1 < n)
                {   // root 涂色
                    m[root][n1+1] = max(m[root][n1+1], v1+root->val);
                }

            }
        }
        else
        {   // root 不涂色，root相连的有色节点为0
            m[root][0] = max(m[root][0], 0);
            // root 涂色
            m[root][1] = max(m[root][1], root->val);
        }
    }
};

880 ms 249.1 MB C++

---

我的CSDN博客地址 https://michael.blog.csdn.net/

长按或扫码关注我的公众号（Michael阿明），一起加油、一起学习进步！