9 - 树形DP

非常哈人动态规划，使我脑子旋转。

概念

树形 DP，即在树上进行的 DP。由于树固有的递归性质，树形 DP 一般都是递归进行的。

例题

没有上司的舞会

285. 没有上司的舞会 - AcWing题库

状态表示：

f[u, 0]，所有从以u为根节点的子树中选择，并且不选u的方案。

f[u, 1]，所有从以u为根节点的子树中选择，并且选择u的方案。

我们要求它们的最大值。

状态计算：

先递归处理所有子树S1, S2, ...，算出f[S1, 0], f[S1, 1], ...，然后才能得出：

• f[u, 0] += max(f[S1, 0], f[S1, 1]) + max(......)。
• f[u, 1] += f[S1, 1] + ...。

参考代码如下：

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

struct Graph
{
    vector<int> h, happy, hasFather, e, ne;
    int idx;
    
    void add(int a, int b)
    {
        hasFather[b] = 1;
        e[idx] = b;
        ne[idx] = h[a];
        h[a] = idx++;
    }
    
    Graph(int n) : h(n + 1, -1), happy(n + 1, 0), hasFather(n + 1, 0),
                   e(2 * n + 5), ne(2 * n + 5), idx(0) {}
};

void dfs(vector<vector<int>>& f, Graph& g, int root)
{
    f[root][1] = g.happy[root];
    
    for (int idx = g.h[root]; idx != -1; idx = g.ne[idx])
    {
        int son = g.e[idx];
        dfs(f, g, son);
        
        f[root][0] += max(f[son][0], f[son][1]);
        f[root][1] += f[son][0];
    }
}

int main()
{
    int n;
    cin >> n;
    
    Graph g(n);
    for (int i = 1; i <= n; ++i)
        cin >> g.happy[i];    
    for (int i = 1; i <= n - 1; ++i)
    {
        int a, b;
        cin >> a >> b;
        g.add(b, a);
    }
    // 找根节点
    int root = 1;
    while (g.hasFather[root])
        ++root;
        
    // 通过DFS进行树形dp
    vector<vector<int>> dp(n + 1, vector<int>(2, 0));
    dfs(dp, g, root);
    
    cout << max(dp[root][0], dp[root][1]);
}

练习题

打家劫舍3

int rob(TreeNode* root) {
// 树形DP
// dp(i, 0): 所有从以 i 为根节点的⼦树中选择，并且不选 i 的方案。
// dp(i, 1): 所有从以 i 为根节点的⼦树中选择，并且选择 i 的方案。
// 求价格最大值

    auto Traverse = [](auto&& Traverse, TreeNode* node) -> vector<int>
    {
        if (node == nullptr) return {0, 0};

        vector<int> leftRes = Traverse(Traverse, node->left);
        vector<int> rightRes = Traverse(Traverse, node->right);

        vector<int> result(2);

        // 不偷该节点, 那么在左右节点中选择最大结果
        result[0] = max(leftRes[0], leftRes[1]) +
                    max(rightRes[0], rightRes[1]);

        // 偷该节点, 那么不能偷左右节点 
        result[1] = node->val + leftRes[0] + rightRes[0];
        return result;
    };

    vector<int> result = Traverse(Traverse, root);
    return max(result[0], result[1]);
}