每日一题 337. 打家劫舍 III

337. 打家劫舍 III

背包问题的变形

class Solution {
public:
    int rob(TreeNode* root) {
        unordered_map<TreeNode*,int> mpYes;
        unordered_map<TreeNode*,int> mpNo;

        int ans = 0;
         dfs(root,mpYes,mpNo);
        return max(mpNo[root],mpYes[root]);
    }

    void dfs(TreeNode* root,unordered_map<TreeNode*,int> &mpYes, unordered_map<TreeNode*,int> &mpNo )
    {
        if(!root)
        {
            mpYes[root] = 0;
            mpNo[root]  = 0;
            return;
        }
        if(!root->left && !root->right)
        {
            mpYes[root] = root->val;
            mpNo[root]  = 0;
            return ; 
        }
        if(!mpYes.count(root->left)){
            dfs(root->left,mpYes,mpNo);
        } 
        if(!mpYes.count(root->right))
        {
            dfs(root->right,mpYes,mpNo);
        }
        
        int yesVal = root->val;
        int noVal  = 0;
        yesVal += mpNo[root->left] + mpNo[root->right];
        noVal  += (max(mpNo[root->left]+mpNo[root->right],mpYes[root->left]+mpYes[root->right]));
        noVal = max(noVal,mpNo[root->left]+mpYes[root->right]);
        noVal = max(noVal,mpNo[root->right]+mpYes[root->left]);
        mpYes[root] = yesVal;
        mpNo[root]  = noVal;
        //cout<<root->val<<" "<<yesVal<<" "<<noVal<<endl;

    }
};

更高效的算法

int rob(TreeNode* root) {
        pair<int, int> result = dfs(root);
        return max(result.first, result.second);
    }

private:
    // 返回值为 pair<int, int>，分别表示 (不偷该节点的最大值, 偷该节点的最大值)
    pair<int, int> dfs(TreeNode* root) {
        if (!root) return {0, 0};

        // 后序遍历：先计算左右子树
        pair<int, int> left = dfs(root->left);
        pair<int, int> right = dfs(root->right);

        // 不偷当前节点
        int notRob = max(left.first, left.second) + max(right.first, right.second);
        // 偷当前节点
        int rob = root->val + left.first + right.first;

        return {notRob, rob};
    }