オットセイの経営日誌

データサイエンス系ベンチャーを経営してます。経営のこと、趣味のことつぶやきます。

LeetCode / House Robber

プログラミングの技術ではなく、高校(中学?)数学の知識を頭の片隅から引っ張り出す必要がある問題。

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

You are a professional robber planning to rob houses along a street. Each house has a certain amount of money stashed, the only constraint stopping you from robbing each of them is that adjacent houses have security system connected and it will automatically contact the police if two adjacent houses were broken into on the same night.

Given a list of non-negative integers representing the amount of money of each house, determine the maximum amount of money you can rob tonight without alerting the police.

Example 1:

Input: [1,2,3,1]
Output: 4
Explanation: Rob house 1 (money = 1) and then rob house 3 (money = 3).
Total amount you can rob = 1 + 3 = 4.
Example 2:

Input: [2,7,9,3,1]
Output: 12
Explanation: Rob house 1 (money = 2), rob house 3 (money = 9) and rob house 5 (money = 1).
Total amount you can rob = 2 + 9 + 1 = 12.

道路に家が立ち並んでいて、リストの各値が各家が貯めているお金だとして、泥棒が各家からお金を盗むケースを想定します。

ただし、隣り合う家から盗ってしまうと警察に連絡されてしまうという前提で、最大限盗めるお金の額を返す問題です。

解答・解説

誤答1

一応示しておくと、以下のようにやってしまうと間違いです。

class Solution:
    def rob(self, nums: List[int]) -> int:
        return max(sum(nums[::2]), sum(nums[1::2]))

初めの要素から1つ飛ばしで和をとる場合と、2つ目の要素から1つ飛ばしで和をとる場合を比べてmaxをとっていますが、

これだと例えば[2,1,1,2]の場合に正しい答えを返せません。

解法1

リストの各要素についてループを回していき、「2つ前の要素までの和の最大値に今の要素を足した値」と「1つ前の要素までの和の最大値」を比べてmaxをとっていく、という方法をとります。
数学的には漸化式と呼ばれる関係になりますが、コードで書くと以下のようになることを意味します。

nums[0] = nums[0]
nums[1] = max(nums[0], nums[1])
nums[k] = max(k + nums[k-2], nums[k-1])

以上の考え方を用いると、以下のように書けます。
numsと同じ長さのリストdpを用意し、漸化式に基づいて計算したnumsの各要素までの最大値を格納します。最終的に返すべき値はdpの最後の要素になります。

class Solution:
    def rob(self, nums: List[int]) -> int:
        
        if not nums: return 0
        if len(nums) == 1: return nums[0]
        
        dp = [0] * len(nums)
        dp[0] = nums[0]
        dp[1] = max(nums[0], nums[1])
        for i in range(2, len(nums)):
            dp[i] = max(nums[i] + dp[i-2], dp[i-1])
            
        return dp[-1]

numsと同じ長さのリストを用意したので、時間計算量、空間計算量ともにO(n)となります。

解法2

以下の書き方で、空間計算量をO(1)に抑えることができます。

class Solution:
    def rob(self, nums: List[int]) -> int:
        nlast, nmax = 0, 0
        for i in nums: 
            nlast, nmax = nmax, max(nlast + i, nmax)
        return nmax