House Robber III

c++

/**
* Definition for a binary tree node.
* struct TreeNode {
*     int val;
*     TreeNode *left;
*     TreeNode *right;
*     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class Solution {
public:
int rob(TreeNode* root) {
//left->right->root
if (root == nullptr) return 0;
int we_rob = 0, we_not_rob = 0;
tryRob(root, we_rob, we_not_rob);
return max(we_rob, we_not_rob);
}
private:
void tryRob(const TreeNode* root, int& we_rob, int& we_not_rob) {
if (root->left == nullptr && root->right == nullptr) {
we_rob = root->val;
we_not_rob = 0;
return;
}
int cur_rob_left = 0;
int cur_not_rob_left = 0;
int cur_rob_right = 0;
int cur_not_rob_right = 0;
if(root->left)
tryRob(root->left,  cur_rob_left,  cur_not_rob_left);
if(root->right)
tryRob(root->right, cur_rob_right, cur_not_rob_right);
we_rob = cur_not_rob_left + cur_not_rob_right + root->val;
int tmp1 = max(cur_rob_left + cur_rob_right, cur_not_rob_left + cur_not_rob_right);
int tmp2 = max(cur_not_rob_left + cur_rob_right, cur_rob_left + cur_not_rob_right);
we_not_rob = max(tmp1, tmp2);
}
};

python

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
def rob(self, root):
"""
:type root: TreeNode
:rtype: int
"""
if not root: return 0
we_rob, we_not_rob = self.tryRob(root)
return max(we_rob, we_not_rob)

def tryRob(self, root):
if not root.left and not root.right:
return root.val, 0
cur_rob_left,  cur_not_rob_left = 0, 0
cur_rob_right, cur_not_rob_right = 0, 0

if root.left:
cur_rob_left,  cur_not_rob_left = self.tryRob(root.left)
if root.right:
cur_rob_right, cur_not_rob_right = self.tryRob(root.right)

we_rob = cur_not_rob_left + cur_not_rob_right + root.val
we_not_rob = max(cur_rob_left + cur_rob_right,
cur_not_rob_left + cur_not_rob_right,
cur_not_rob_left + cur_rob_right,
cur_rob_left + cur_not_rob_right)
return we_rob, we_not_rob