思路:如果选择树的当前节点就不选择左右孩子节点,为了避免对节点重复dfs超时,使用map记录已经访问过的节点。
代码如下:
class Solution {
public:
map<TreeNode *, int> mp;
int dfs(TreeNode* tree, int type) {
if (tree == NULL) return 0;
if (mp[tree] != 0) return mp[tree];
if (type == 1) mp[tree] = dfs(tree->left, 0) + dfs(tree->right, 0);
else mp[tree] = max(tree->val + dfs(tree->left, 1) + dfs(tree->right, 1), dfs(tree->left, 0) + dfs(tree->right, 0));
return mp[tree];
}
int rob(TreeNode* root) {
int ans = dfs(root, 0);
return ans;
}
};
为了区分是否选择了当前节点,在dfs中使用参数type记录,type==1说明当前无法选择,type==0说明可以选择访问当前节点或者不访问当前节点,取max。
结果还是WA,仔细思考后发现,type限制了当前节点的map值,如果当前节点可选可不选,代码的max操作先type==1,如果mp值更新为不为0的数,之后的type==0遍历到直接取map值,这样导致了WA。
所以在记忆化搜索的时候要注意对记录的值的定义,在设计dfs的时候思考自己求得的是否是自己定义的。
AC代码:
class Solution {
public:
map<TreeNode *, int> mp;
int dfs(TreeNode* tree) {
if (tree == NULL) return 0;
if (mp[tree] != 0) return mp[tree];
int a = tree->val;
if (tree->left){
a += dfs(tree->left->left) + dfs(tree->left->right);
}
if (tree->right){
a += dfs(tree->right->left) + dfs(tree->right->right);
}
int b = dfs(tree->left) + dfs(tree->right);
mp[tree] = max(a, b);
return mp[tree];
}
int rob(TreeNode* root) {
int ans = dfs(root);
return ans;
}
};