코테 준비

[프로그래머스/DP/python] 도둑질

우디혜 2021. 2. 16. 01:51

 

처음에는 recursion을 이용하려 했지만 시간 초과가 날 것 같아서 다른 방식을 찾아봐야했다.

결국 스터디에서 친구의 코드를 참고해서 문제를 풀게 되었다.

 

DP 문제는 보통

dp[n] = dp[n -1] + dp[n - 2]

dp[n] = max(dp[n - 1], dp[n - 2])

...

이런 느낌의 점화식으로 풀 수 있는 문제인데 이번 문제도 그랬다.

 

로직

[10, 20, 30, 40, 50, 60] 이 있으면

인덱스가 2인 위치에서 최대가 될 수 있는 경우는 max(30 + 10, 20) -> max 값을 3의 위치에 업데이트
[10, 20, 40, 40, 50, 60] 

인덱스가 3인 위치에서 최대가 될 수 있는 경우는 max(40 + 20, 40) -> max 값을 4의 위치에 업데이트

[10, 20, 40, 60, 50, 60] 

인덱스가 4인 위치에서 최대가 될 수 있는 경우는 max(50 + 40, 60) -> max 값을 5의 위치에 업데이트

...

 

이런 느낌으로 진행하면 된다.

결국 dp[n] = max(dp[n - 1], dp[n - 2] + dp[n]) 점화식을 사용하면 각 위치에서의 최댓값을 구할 수 있다.

 

주의할 점이 마지막 집과 처음 집을 동시에 선택하는 경우를 피해야하기 때문에, 첫 번째 집을 선택하는 경우와 선택하지 않는 경우를 나눠야 한다.

코드

import copy
def solution(money):
    memo = money.copy()
    if len(money) == 3:
        return max(money[0], money[1], money[2])
    
    memo[1] = max(memo[0], memo[1])
    for i in range(2, len(memo) - 1):
        memo[i] = max(memo[i - 2] + memo[i], memo[i - 1])
    
    money[0] = 0
    for i in range(2, len(money)):
        money[i] = max(money[i - 2] + money[i], money[i - 1])
    
    first_max, second_max = max(memo), max(money)
    
    return max(first_max, second_max)

 

솔직히 예쁜 코드는 아니다. 그래서 다른 분의 코드를 참고해서 리펙토링해본 결과

 

def solution(money):
    
    if len(money) == 3:
        return max(money[0], money[1], money[2])
    
    # 첫 번째 집을 방문하는 경우
    include_first = [money[0]]
    include_first.append(max(money[0], money[1]))
                      
    for i in range(2, len(money) - 1):
        include_first.append(max(include_first[i - 2] + money[i], include_first[i - 1]))
    
    # 첫 번째 집을 방문하지 않는 경우
    exclude_first = [0, money[1]]
                      
    for i in range(2, len(money)):
        exclude_first.append(max(exclude_first[i - 2] + money[i], exclude_first[i - 1]))
    
    return max(include_first[-1], exclude_first[-1])