337. 打家劫舍 III——python3

题目描述

在上次打劫完一条街道之后和一圈房屋后，小偷又发现了一个新的可行窃的地区。这个地区只有一个入口，我们称之为“根”。除了“根”之外，每栋房子有且只有一个“父”房子与之相连。一番侦察之后，聪明的小偷意识到“这个地方的所有房屋的排列类似于一棵二叉树”。 如果两个直接相连的房子在同一天晚上被打劫，房屋将自动报警。

计算在不触动警报的情况下，小偷一晚能够盗取的最高金额。

示例1：
输入：[3,2,3,null,3,null,1]

     3
    / \
   2   3
    \   \ 
     3   1

输出：7
解释：小偷一晚能够盗取的最高金额 = 3 + 3 + 1 = 7

示例2：
输入：[3,4,5,1,3,null,1]

     3
    / \
   4   5
  / \   \ 
 1   3   1

输出：9
解释：小偷一晚能够盗取的最高金额 = 4 + 5 = 9

示例代码

house_robber_iii.py

# -*- coding: utf-8 -*-

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

class Solution():
    """给出打家劫舍III的答案"""

    def money_rob_or_norob(self, root):
        """
        给出偷当前节点和不偷当前节点，分别能得到的最大金额
        例如，对于根节点的左右孩子节点
        分别定义元组my_left和my_right，其中：
            [0]：表示如果偷当前节点，可以得到的最大金额
            [1]：表示如果不偷当前节点，可以得到的最大金额
        则最终的最大金额为以下5个数值中的最大值
            偷根节点
                1. 不偷左右孩子：root.val + m_l[1] + m_r[1]
            不偷根节点
                2. 偷左，偷右：m_l[0] + m_r[0]
                3. 偷左，不偷右：m_l[0] + m_r[1]
                4. 不偷左，偷右：m_l[1] + m_r[0]
                5. 不偷左，不偷右：m_l[1] + m_r[1]
        以此类推。。。
        """
        # 若root为空，返回0
        if root is None:
            return (0, 0)

        my_left = self.money_rob_or_norob(root.left)
        my_right = self.money_rob_or_norob(root.right)

        rob_current = root.val + my_left[1] + my_right[1]
        norob_current = max(
            max(my_left[0] + my_right[0], my_left[0] + my_right[1]),
            max(my_left[1] + my_right[0], my_left[1] + my_right[1]))

        return (rob_current, norob_current)

    def rob(self, root):
        """
        给出能偷取的最大金额
        :type root: TreeNode
        :rtype: int
        """
        return max(self.money_rob_or_norob(root))

test_house_robber_iii.py

# -*- coding: utf-8 -*-

import unittest
from house_robber_iii import Solution

class TreeNode():
    """表示二叉树的类"""

    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class TestAsteroidCollision(unittest.TestCase):
    """针对Solution类的测试"""

    def setUp(self):
        """创建一个Solution类，供使用的测试方法使用"""
        self.my_solution = Solution()

    def test_empty_tree(self):
        """空树"""
        root = None
        result = self.my_solution.rob(root)
        expected_result = 0
        self.assertEqual(result, expected_result)

    def test_only_one_node_value_0(self):
        """
        只有一个根节点
        而且根节点为0
        """
        root = TreeNode(0)
        result = self.my_solution.rob(root)
        expected_result = 0
        self.assertEqual(result, expected_result)

    def test_only_one_node_value_1(self):
        """
        只有一个根节点
        而且根节点为1
        """
        root = TreeNode(1)
        result = self.my_solution.rob(root)
        expected_result = 1
        self.assertEqual(result, expected_result)

    def test_rob_root_norob_son(self):
        """
        偷根节点
        不偷左右孩子
        """
        root = TreeNode(10)
        root.left = TreeNode(3)
        root.right = TreeNode(4)
        result = self.my_solution.rob(root)
        expected_result = 10
        self.assertEqual(result, expected_result)

    def test_rob_root_rob_grandson(self):
        """
        偷根节点
        偷孙节点
        """
        root = TreeNode(2)
        root.left = TreeNode(3)
        root.right = TreeNode(4)
        root.left.left = TreeNode(2)
        root.right.right = TreeNode(4)
        result = self.my_solution.rob(root)
        expected_result = 8
        self.assertEqual(result, expected_result)

    def test_rob_left_norob_right_rob_rightson(self):
        """
        偷左孩子
        不偷右孩子，但是偷右孩子的孩子
        """
        root = TreeNode(1)
        root.left = TreeNode(10)
        root.right = TreeNode(2)
        root.left.left = TreeNode(3)
        root.left.right = TreeNode(4)
        root.right.left = TreeNode(7)
        root.right.right = TreeNode(8)
        result = self.my_solution.rob(root)
        expected_result = 25
        self.assertEqual(result, expected_result)

unittest.main()