树形动态规划

MrHe··4 min read

树状DP,是一种在树形结构上进行动态规划的方法。它利用树的递归性质,从叶子节点向上传递信息,最终得到根节点的解。今天我来聊聊树状DP的几种经典应用和解题思路。

树状DP的基本思想

树状DP的核心思想是利用树的递归结构,通过后序遍历(或类似的遍历方式)自底向上地处理每个节点,将子节点的信息整合到父节点中。通常,我们会定义一个状态数组dp[u]表示以节点u为根的子树的某个属性值。

例1:没有上司的舞会

这是一个经典的树状DP问题。问题描述:有n个员工,每个员工有一个快乐值,现在要举办舞会,要求没有直接上下级关系的员工同时参加。求最大快乐值。

解法一:基础树状DP

定义dp[u][0]表示不选u节点时,以u为根的子树的最大快乐值;dp[u][1]表示选u节点时,以u为根的子树的最大快乐值。

状态转移方程:

  • 如果不选u节点,那么它的子节点可选可不选:dp[u][0] = sum(max(dp[v][0], dp[v][1])),其中vu的子节点。

  • 如果选u节点,那么它的子节点都不能选:dp[u][1] = happy[u] + sum(dp[v][0]),其中vu的子节点。

void dfs(int u, int father) {
    dp[u][0] = 0;
    dp[u][1] = happy[u];

    for (int v : children[u]) {
        if (v == father) continue;
        dfs(v, u);
        dp[u][0] += Math.max(dp[v][0], dp[v][1]);
        dp[u][1] += dp[v][0];
    }
}

解法二:记忆化搜索

对于一些特殊的树结构,比如二叉树,我们可以用记忆化搜索的方式来实现树状DP。

int[][] dfs(TreeNode u) {
    if (u == null) return new int[]{0, 0};

    int[] left = dfs(u.left);
    int[] right = dfs(u.right);

    int select = u.val + left[0] + right[0];
    int notSelect = Math.max(left[0], left[1]) + Math.max(right[0], right[1]);

    return new int[]{notSelect, select};
}

解法三:BFS+拓扑排序

对于一些特殊的树结构,比如链式树,我们可以用BFS+拓扑排序的方式来实现树状DP。

void treeDP() {
    Queue<Integer> queue = new LinkedList<>();
    int[] inDegree = new int[n];

    // 初始化入度
    for (int i = 0; i < n; i++) {
        for (int v : children[i]) {
            inDegree[v]++;
        }
    }

    // 将所有叶子节点入队
    for (int i = 0; i < n; i++) {
        if (inDegree[i] == 1) { // 叶子节点
            queue.offer(i);
        }
    }

    while (!queue.isEmpty()) {
        int u = queue.poll();

        // 更新dp[u]
        dp[u][0] = 0;
        dp[u][1] = happy[u];

        for (int v : children[u]) {
            dp[u][0] += Math.max(dp[v][0], dp[v][1]);
            dp[u][1] += dp[v][0];

            if (--inDegree[v] == 1) {
                queue.offer(v);
            }
        }
    }
}

例2:二叉树的最大路径和

问题描述:给定一个二叉树,找出其最大路径和。路径可以是从任意节点到任意节点,但必须经过父节点关系。

解法一:递归树状DP

定义dp[u]表示以u为根的子树中,包含u节点的最大路径和。

int maxPathSum(TreeNode root) {
    maxSum = Integer.MIN_VALUE;
    dfs(root);
    return maxSum;
}

int dfs(TreeNode u) {
    if (u == null) return 0;

    int left = Math.max(0, dfs(u.left)); // 如果子路径和为负,则不选
    int right = Math.max(0, dfs(u.right));

    // 计算包含当前节点的最大路径和
    maxSum = Math.max(maxSum, u.val + left + right);

    // 返回包含当前节点且只能选择一条子路径的最大和
    return u.val + Math.max(left, right);
}

解法二:后序遍历+全局变量

这种解法与解法一类似,但更强调遍历顺序。

int maxPathSum(TreeNode root) {
    maxSum = Integer.MIN_VALUE;
    postOrder(root);
    return maxSum;
}

int postOrder(TreeNode u) {
    if (u == null) return 0;

    int left = Math.max(0, postOrder(u.left));
    int right = Math.max(0, postOrder(u.right));

    maxSum = Math.max(maxSum, u.val + left + right);

    return u.val + Math.max(left, right);
}

解法三:分治法

将问题分解为左子树、右子树和当前节点三个部分。

class Result {
    int maxPath; // 包含当前节点的最大路径和
    int maxBranch; // 从当前节点向下的最大分支和

    public Result(int maxPath, int maxBranch) {
        this.maxPath = maxPath;
        this.maxBranch = maxBranch;
    }
}

Result divideConquer(TreeNode u) {
    if (u == null) return new Result(Integer.MIN_VALUE, 0);

    Result left = divideConquer(u.left);
    Result right = divideConquer(u.right);

    // 计算当前节点的最大分支和
    int maxBranch = u.val + Math.max(0, Math.max(left.maxBranch, right.maxBranch));

    // 计算包含当前节点的最大路径和
    int maxPath = u.val;
    if (left.maxBranch > 0) maxPath += left.maxBranch;
    if (right.maxBranch > 0) maxPath += right.maxBranch;

    // 与左右子树的最大路径和比较
    maxPath = Math.max(maxPath, Math.max(left.maxPath, right.maxPath));

    return new Result(maxPath, maxBranch);
}

例3:树的直径

问题描述:给定一棵树,求树中最长路径的长度(路径上边的数量)。

解法一:两次DFS

第一次DFS从任意节点出发,找到距离它最远的节点u;第二次DFS从u出发,找到距离u最远的节点v,则uv的路径就是树的直径。

int[] dfs(int u, int father) {
    int max1 = 0, max2 = 0; // 最大和次大距离
    for (int v : children[u]) {
        if (v == father) continue;
        int d = dfs(v, u)[0] + 1;
        if (d > max1) {
            max2 = max1;
            max1 = d;
        } else if (d > max2) {
            max2 = d;
        }
    }
    diameter = Math.max(diameter, max1 + max2);
    return new int[]{max1, u};
}

int treeDiameter(int[][] edges) {
    // 构建邻接表
    List<Integer>[] adj = new List[n];
    for (int i = 0; i < n; i++) {
        adj[i] = new ArrayList<>();
    }
    for (int[] edge : edges) {
        adj[edge[0]].add(edge[1]);
        adj[edge[1]].add(edge[0]);
    }

    diameter = 0;
    dfs(0, -1);
    return diameter;
}

解法二:树状DP

定义dp[u]表示以u为根的子树中,从u出发向下的最长路径长度。

int[] dp = new int[n];
int diameter = 0;

void dfs(int u, int father) {
    int max1 = 0, max2 = 0; // 最大和次大距离
    for (int v : adj[u]) {
        if (v == father) continue;
        dfs(v, u);
        if (dp[v] + 1 > max1) {
            max2 = max1;
            max1 = dp[v] + 1;
        } else if (dp[v] + 1 > max2) {
            max2 = dp[v] + 1;
        }
    }
    dp[u] = max1;
    diameter = Math.max(diameter, max1 + max2);
}

解法三:BFS+动态规划

对于一些特殊的树结构,比如链式树,我们可以用BFS+动态规划的方式来实现。

int treeDiameter(int[][] edges) {
    // 构建邻接表
    List<Integer>[] adj = new List[n];
    for (int i = 0; i < n; i++) {
        adj[i] = new ArrayList<>();
    }
    for (int[] edge : edges) {
        adj[edge[0]].add(edge[1]);
        adj[edge[1]].add(edge[0]);
    }

    int[] dp = new int[n];
    int diameter = 0;
    Queue<Integer> queue = new LinkedList<>();
    int[] inDegree = new int[n];

    // 初始化入度
    for (int i = 0; i < n; i++) {
        inDegree[i] = adj[i].size();
        if (inDegree[i] == 1) { // 叶子节点
            queue.offer(i);
        }
    }

    while (!queue.isEmpty()) {
        int u = queue.poll();

        for (int v : adj[u]) {
            if (--inDegree[v] == 1) {
                queue.offer(v);
            }
            dp[v] = Math.max(dp[v], dp[u] + 1);
            diameter = Math.max(diameter, dp[v]);
        }
    }

    return diameter;
}

总结

树状DP是一种强大的算法思想,它利用树的递归性质,通过定义合适的状态和状态转移方程,可以高效地解决许多树形结构上的问题。在实际应用中,我们可以根据问题的特点选择不同的实现方式,如递归、记忆化搜索、BFS+拓扑排序等。掌握树状DP的关键在于理解问题的本质,正确定义状态,并找到合适的状态转移方程。