Leetcode - 337. House Robber III

문제

이진 트리 구조로 연결된 집이 존재한다. 인접한 두 집이 털렸을때 경보가 울린다. 경보를 울리지 않고 모든 집을 털때 가장 많이 훔칠 수 있는 금액은?

Input: root = [3,2,3,null,3,null,1]
Output: 7
Explanation: Maximum amount of money the thief can rob = 3 + 3 + 1 = 7.

Input: root = [3,4,5,1,3,null,1]
Output: 9
Explanation: Maximum amount of money the thief can rob = 4 + 5 = 9.

해결 아이디어

• 인접한 두 노드를 훔치지 않아야한다. 이전 House Robber 문제들 처럼 현재 노드에서는 훔치거나, 훔치지 않는 두가지 경우가 존재. 이에 따라 점화식을 세우고 dfs순회를 하면된다.
• 나이브 하게 dfs를 하면 아래와 같다. -> 결과는 TLE

class Solution {
    int dfs(TreeNode* root, bool visit) {
        if (!root)
            return 0;

        if (visit) { // if (visit == true)  can add self node, or not add 
            int self_added = dfs(root->left, false) + dfs(root->right, false) + root->val;
            int self_passed = dfs(root->left, true) + dfs(root->right, true);
            return max(self_added, self_passed);
        } else { // must not add self node
            return dfs(root->left, true) + dfs(root->right, true);
        }
    }
public:
    int rob(TreeNode* root) {
        return dfs(root, true);
    }
};

풀이

위 풀이에 memoization을 추가하면 아래와 같다. 주의할 점은 memo 해시테이블을 두개를 두어야 한다는점이다(힌트로 알게된 내용). 현재 노드에서 훔칠수도 있고 안훔칠수도 있기때문에 그 두가지 경우에따라 현재 노드에 max값이 달라진다. 따라서 두 종류의 값을 저장해야한다.

class Solution {
public:
    unordered_map<TreeNode *, int> with_mem;
    unordered_map<TreeNode *, int> without_mem;
    
    int dfs(TreeNode* root, bool visit) {
        if (!root)
            return 0;
        
        if (visit) {
            if (with_mem.find(root) != with_mem.end())
                return with_mem[root];
            int with_self = dfs(root->left, false) + dfs(root->right, false) + root->val;
            int without_self = dfs(root->left, true) + dfs(root->right, true);
            with_mem[root] = max(with_self, without_self);
            return with_mem[root];
        } else {
            if (without_mem.find(root) != without_mem.end())
                return without_mem[root];
            without_mem[root] = dfs(root->left, true) + dfs(root->right, true);
            return without_mem[root];
        }
    }
    int rob(TreeNode* root) {
        return dfs(root, true);
    }
};