Leetcode 337. 打家劫舍 III 的优化过程

69 阅读1分钟

首先是暴力dfs

image.png

class Solution {
    public int rob(TreeNode root) {
        int ans1 = 0;
        int ans2 = 0;
        // root被偷
        ans1 = root.val + dfs(root.left, true) + dfs(root.right, true);
        // root没被偷
        ans2 = dfs(root.left, false) + dfs(root.right, false);
        return Math.max(ans1, ans2);
    }

    // 需要知道: 其父节点有没有被偷
    private int dfs(TreeNode node, boolean stolen) {
        if (node == null) return 0;
        if (stolen) {
            // 其父节点被偷,当前节点不能偷
            return dfs(node.left, false) + dfs(node.right, false);
        } else {
            // 其父节点没被偷,当前节点可以偷,也可以不偷
            return Math.max(dfs(node.left, true) + dfs(node.right, true) + node.val,
                            dfs(node.left, false) + dfs(node.right, false)) ;
        }
    }
}

优化:采用后序遍历的方式,可以将偷没偷的中间结果保存起来,采用DP的方式去解

image.png

class Solution {
    public int rob(TreeNode root) {
        int[] ans = dfs(root);
        return Math.max(ans[0], ans[1]);
    }

    private int[] dfs(TreeNode node) {
        int[] tmp = new int[2];  // tmp[0]: stole; tmp[1]: no stole
        if (node == null)
            return tmp;
        
        int[] left = dfs(node.left);
        int[] right = dfs(node.right);

        // 当前节点
        // no stole, 左右孩子可偷可不偷
        tmp[0] = Math.max(left[0], left[1]) + Math.max(right[0], right[1]);
        // stole, 左右孩子都不能偷
        tmp[1] = node.val + left[0] + right[0]; 

        return tmp;
    }
}