[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);
    }
};