Files
leetcode-master/problems/kamacoder/0131.小美的树上染色.md
programmercarl 26a5b0cc21 update
2024-06-24 17:13:15 +08:00

4.2 KiB
Raw Blame History

131. 小美的树上染色

本题为树形dp 稍有难度,主要在于 递推公式上。

dp数组的定义

dp[cur][1] :当前节点染色,那么当前节点为根节点及其左右子节点中,可以染色的最大数量

dp[cur][0] :当前节点不染色,那么当前节点为根节点及其左右子节点中,可以染色的最大数量

关于 dp转移方程

1、 情况一:

如果当前节点不染色,那就去 子节点 染色 或者 不染色的最大值。

dp[cur][0] += max(dp[child][0], dp[child][1]);

2、情况二

那么当前节点染色的话,这种情况就不好想了。

首先这不是二叉树,每一个节点都有可能 会有n个子节点。

所以我们要分别讨论,每一个子节点的情况 对父节点的影响。

那么父节点 针对每种情况,就要去 最大值, 也就是 dp[cur][1] = max(dp[cur][1], 每个自孩子的情况)

如图假如节点1 是我们要计算的父节点节点2是我们这次要计算的子节点。

选中一个节点2 作为我们这次计算的子节点,父节点染色的话,子节点必染色。

接下来就是计算 父节点1和该子节点2染色的话 以子节点2 为根的 染色节点的最大数量 。

节点2不染色 且 以节点2为根节点的最大 染色数量 + 2 + 2 是因为 节点 1 和 节点2 要颜色了,染色节点增加两个。

代码:dp[child][0] + 2

细心的录友会发现,那我们只计算了 红色框里面的,那么框外 最大的染色数量是多少呢?

先看 作为子节点的节点2 为根节点的最大染色数量是多少? 取一个最值,即 节点2染色 或者 不染色取最大值。

代码:max(dp[child][0], dp[child][1])

那么红框以外的 染色最大节点数量 就是 dp[cur][0] - max(dp[child][0], dp[child][1])

cur是节点1child是节点2

红框以外的染色最大数量 + 父节点1和该子节点2染色的话 以子节点2 为根的 染色节点的最大数量 就是 节点1 染色的最大节点数量。

代码:

dp[cur][1] = max(dp[cur][1], dp[cur][0] - max(dp[child][0], dp[child][1]) + dp[child][0] + 2);

整体代码如下:


#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>
#include <list>

using namespace std;

int maxN = 10005;
vector<vector<int>> dp (maxN, vector<int>(2, 0));
vector<list<int>> grid(maxN); // 邻接表
vector<long> value(maxN); // 存储每个节点的权值


// 在树上进行动态规划的函数
void dpOnTheTree(int cur) {

    for (int child : grid[cur]) {
        // 后序遍历,从下向上计算
        dpOnTheTree(child);
        // 情况一
        dp[cur][0] += max(dp[child][0], dp[child][1]);

    }

    // 计算dp[1] - 当前节点染色
    for (int child : grid[cur]) {
        long mul = value[cur] * value[child]; // 当前节点和相邻节点权值的乘积
        long sqrtNum = (long) sqrt(mul);

        if (sqrtNum * sqrtNum == mul) { // 如果乘积是完全平方数
            // 情况二
            // dp[cur][0] 表示所有子节点 染色或者不染色的 最大染色数量
            // max(dp[child][0], dp[child][1])  需要染色节点的孩子节点的最大染色数量
            // dp[cur][0] - max(dp[child][0], dp[child][1]) 除了要染色的节点及其子节点,其他孩子的最大染色数量
            // 最后 + dp[child][0] + 2  就是本节点染色的最大染色节点数量
            dp[cur][1] = max(dp[cur][1], dp[cur][0] - max(dp[child][0], dp[child][1]) + dp[child][0] + 2);
        }
    }

}

int main() {

    int n;
    cin >> n; // 输入节点数量

    // 读取节点权值
    for (int i = 1; i <= n; ++i) {
        cin >> value[i];
    }

    // 构建树的邻接表
    for (int i = 1; i < n; ++i) {
        int x, y;
        cin >> x >> y;
        grid[x].push_back(y);
    }

    // 从根节点节点1开始进行动态规划
     dpOnTheTree(1);

    // 输出最大染色节点数量
    cout << max(dp[1][0], dp[1][1]) << endl;

    return 0;
}