[LeetCode] 337. House Robber III

題目鏈接: https://leetcode.com/problems/house-robber-iii/description/

Description

The thief has found himself a new place for his thievery again. There is only one entrance to this area, called the “root.” Besides the root, each house has one and only one parent house. After a tour, the smart thief realized that “all houses in this place forms a binary tree”. It will automatically contact the police if two directly-linked houses were broken into on the same night.

Determine the maximum amount of money the thief can rob tonight without alerting the police.

Example 1:

     3
    / \
   2   3
    \   \ 
     3   1

Maximum amount of money the thief can rob = 3 + 3 + 1 = 7.

Example 2:

     3
    / \
   4   5
  / \   \ 
 1   3   1

Maximum amount of money the thief can rob = 4 + 5 = 9.

解題思路

題意轉化爲在二叉樹中，求節點的子集中權值之和的最大值，並且要求子集中所有節點都不相鄰。

對於每一個節點，都有兩種狀態，選或者不選。並且當某節點被選擇時，其父節點和子節點都只有一種狀態，即不能被選；當某節點不被選擇時，其父節點和子節點分別有兩種狀態，選或者不選。因此，需要一個長度爲 2 的數據結構，用以分別保存某節點被選和不被選時，其子樹中滿足要求的權值之和的最大值。在 c++ 中，可以使用 pair<int, int> 來記錄，其中 first 記錄節點被選擇時的最大值，second 記錄節點不被選擇時的最大值。

由於每一個節點在計算其子樹的最大值時，需要使用到左右子樹的最大值，因此可以使用後序遍歷來達到這個目的。

具體步驟爲，後序遍歷二叉樹，對於每一個節點在獲取其左右子樹的最大值後，計算該節點被選擇和不被選擇時整個子樹的最大值。對於當前節點，兩種狀態的最大值分別爲

• 被選擇時：最大值爲當前節點權值加上左右子樹都不被選擇時的最大值；
• 不被選擇時：最大值爲左孩子被選或者不被選兩種狀態中的最大值，加上右孩子被選或者不被選兩種狀態中的最大值。

遞歸下去，最終在根節點取被選或者不被選的最大值即爲結果。

Code

class Solution {
public:
    int rob(TreeNode *root) {
        pair<int, int> res = help(root);
        return max(res.first, res.second);
    }

    pair<int, int> help(TreeNode *root) {
        if (root == NULL) return make_pair(0, 0);
        pair<int, int> left = help(root->left);
        pair<int, int> right = help(root->right);

        int select = root->val + left.second + right.second;
        int noselect = max(left.first, left.second) + max(right.first, right.second);
        return make_pair(select, noselect);
    }
};