LeetCode 337. 打家劫舍 III

◎

本节目标

后序同时返回选当前节点和不选当前节点的最大金额。

题意与约束

二叉树相邻父子不能同时选择，返回可选节点值之和的最大值。

第一反应与重复子问题

选择一个节点会限制子节点，不选择则子节点可各自最优；每棵子树都需要同样的两种答案。

状态定义与转移推导

后序返回 (take, skip)。take = val + left.skip + right.skip，skip = max(left) + max(right)。

正确性依据

父节点的唯一约束只涉及直接子节点，给定父节点选或不选后左右子树独立，二元状态穷尽所有合法情况。

样例执行过程

示例树根为 3 时，选择根会搭配两个子树的 skip，最终取到 7。

代码实现

C++

C++17

#include <algorithm>
#include <functional>
#include <utility>
using namespace std;

class Solution {
public:
    int rob(TreeNode* root) {
        function<pair<int, int>(TreeNode*)> dfs = [&](TreeNode* node) -> pair<int, int> {
            if (!node) {
                return {0, 0};
            }
            auto left = dfs(node->left);
            auto right = dfs(node->right);
            return {
                node->val + left.second + right.second,
                max(left.first, left.second) + max(right.first, right.second),
            };
        };
        auto result = dfs(root);
        return max(result.first, result.second);
    }
};

Python

Python 3

class Solution:
    def rob(self, root):
        if not root:
            return 0

        states = {}
        stack = [(root, False)]
        while stack:
            node, visited = stack.pop()
            if not node:
                continue
            if visited:
                left_take, left_skip = states.get(id(node.left), (0, 0))
                right_take, right_skip = states.get(id(node.right), (0, 0))
                states[id(node)] = (
                    node.val + left_skip + right_skip,
                    max(left_take, left_skip) + max(right_take, right_skip),
                )
            else:
                stack.append((node, True))
                stack.append((node.right, False))
                stack.append((node.left, False))
        return max(states[id(root)])

复杂度分析

每个节点计算一次，时间 O(n)。C++ 实现使用递归，调用栈为 O(h)；Python 实现使用显式栈并保存每个节点的状态，额外空间为 O(n)。

边界与易错点

空树返回零；TreeNode 由平台提供，源码不应重复定义，状态也不能残留到下一次调用。

模式迁移

树上的相邻互斥可用选／不选状态；若约束变为覆盖关系，状态数会增加。回到树形动态规划。