Leetcode 337 Solution

This article provides solution to leetcode question 337 (house-robber-iii)

https://leetcode.com/problems/house-robber-iii

Solution

/**
 * 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 {
    std::map<pair<TreeNode*, bool>, int> cache;

public:
    int rob(TreeNode* root, bool robThisLayer) {
        auto key = make_pair(root, robThisLayer);
        if (cache.find(key) != cache.end())
            return cache[key];

        if (root == NULL)
            return 0;

        if (robThisLayer)
        {
            int m1 = rob(root->left, false);
            int m2 = rob(root->right, false);

            return cache[key] = root->val + m1 + m2;
        }
        else
        {
            int left_m1 = rob(root->left, false);
            int left_m2 = rob(root->left, true);
            int right_m1 = rob(root->right, false);
            int right_m2 = rob(root->right, true);

            return cache[key] = max(left_m1, left_m2) + max(right_m1, right_m2);
        }
    }

    int rob(TreeNode* root) {
        int m1 = rob(root, true);
        int m2 = rob(root, false);

        return max(m1, m2);
    }
};