小美的树上染色

这道题是树形DP类型：

• dp[u][0]：不染色节点u的情况下，染色节点数最大值；
• dp[u][1]：染色节点u的情况下，染色节点数最大值。

然后是状态转移方程：

• dp[u][0] += max(dp[i][0], dp[i][1])，其中i和u在一个图中。不染色u的情况下，染色节点数最大值为它周围子节点的染色节点数最大值之和。
• dp[u][1] = max(dp[u][1], dp[u][0] - max(dp[i][0], dp[i][1]) + dp[i][0] + 2)，其中i和u在一个图中。染色u的情况下，染色节点数最大值需要好好解释。由于题中限制若当前节点染色，那么它相邻的节点只能有一个染色，因此我们需要选择一个相邻节点，使得染色节点数最大。接下来是计算，在节点U不染色且节点I不染色（dp[u][0] - max(dp[i][0], dp[i][1]) + dp[i][0]）的情况下，将两个节点都染色（+2）。

代码如下：

#include <iostream>
#include <vector>
#include <list>
#include <cmath>
using namespace std;

struct Graph
{
    vector<int> value;
    vector<list<int>> g;

    Graph(int n)
    {
        value.resize(n + 1);
        g.resize(n + 1);
    }
};

int main() {
    int n;
    cin >> n;

    Graph graph(n);
    for (int i = 1; i <= n; ++i)
    {
        cin >> graph.value[i];
    }
    for (int i = 0; i < n - 1; ++i)
    {
        int from, to;
        cin >> from >> to;
        graph.g[from].push_back(to);
        graph.g[to].push_back(from);
    }

    // 树形DP:
    // dp[u][0]: 不染红u情况下, 染红节点数最大值
    // dp[u][0] = sum(max(dp[i][0], dp[i][1]), ...), 其中i和u在一个图中
    // dp[u][1]: 染红u情况下, 染红节点数最大值
    // dp[u][1] = max(dp[u][0] - max(dp[i][0], dp[i][1]) + dp[i][0] + 2), 其中i和u在一个图中
    //            U不染色 且                I不染色的情况下,       让U和I染色
    vector<vector<int>> dp(n + 1, vector<int>(2));
    auto DFS = [&](auto&& DFS, int u, int parent) -> void
    {
        // 先求dp[u][0]
        for (int i : graph.g[u])
        {
            // 防止重复求值
            if (parent == i)
            {
                continue;
            }
            DFS(DFS, i, u);
            dp[u][0] += max(dp[i][0], dp[i][1]);
        }

        // 再求dp[u][1]
        for (int i : graph.g[u])
        {
            // 防止重复求值
            if (parent == i)
            {
                continue;
            }

            int root = static_cast<int>(sqrt(graph.value[i] * graph.value[u]));
            if (root * root == graph.value[i] * graph.value[u])
            {
                dp[u][1] = max(dp[u][1], dp[u][0] - max(dp[i][0], dp[i][1]) + dp[i][0] + 2);
            }
        }
    };

    DFS(DFS, 1, 0);
    cout << max(dp[1][0], dp[1][1]);
}
// 64 位输出请用 printf("%lld")

在代码实现中，需要注意的是将子节点和父节点进行比较，避免进行重复计算导致超时，或者结果错误。