Python - No extra space needed


  • 0
    N

    It's not hard to find this transition relation:
    denote the max gain after visiting root as gain(root), then
    gain(root) = max(root.val + gain(root.left.left) + gain(root.left.right) + gain(root.right.left) + gain(root.right.right), gain(root.left) + gain(root.right)

    A subproblem gain(any node) can be met multiple times, therefore we'd better use memorization to speed it up.
    The good news is, we can assume the amount of money in any house is non-negative, namely node.val >= 0 for any node. Then we can modify the val to indicate whether this house is visited or not by changing the non-negative val to non-positive val. Then we know, whenever node.val < 0, node.val * -1 is namely gain(node).
    (of course, if many rooms have zero money, this solution cannot give you much benefit by changing the val sign.)
    Here is the code:

    class Solution(object):
        def rob(self, root):
            """
            :type root: TreeNode
            :rtype: int
            """
            def gain(root):
                if root == None:
                    return 0
                o1 = root.val
                if o1 < 0:
                    return -o1
                if root.left:
                    o1 += gain(root.left.left) + gain(root.left.right)
                if root.right:
                    o1 += gain(root.right.left) + gain(root.right.right)
                o2 = gain(root.left) + gain(root.right)
                o = max(o1,o2)
                root.val = -1*o
                return o
            return gain(root)
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.