LeetCode의 213번째 문제인 House Robber II를 함께 풀어보도록 하겠습니다.
문제
당신은 거리를 따라 집들을 털 계획을 세우고 있는 전문 강도입니다. 각 집에는 특정 금액의 돈이 숨겨져 있습니다. 이 지역의 모든 집들은 원형으로 배치되어 있습니다. 즉, 첫 번째 집은 마지막 집의 이웃입니다. 한편, 인접한 집들은 보안 시스템이 연결되어 있으며, 같은 밤에 두 개의 인접한 집이 털리면 자동으로 경찰에 연락됩니다.
각 집의 돈의 양을 나타내는 정수 배열 nums
가 주어졌을 때, 경찰에 경보를 울리지 않고 오늘밤 털 수 있는 최대 금액을 반환합니다.
예제
입력: nums = [2,3,2]
출력: 3
입력: nums = [1,2,3,1]
출력: 4
입력: nums = [1,2,3]
출력: 3
풀이 1: 재귀
어떻게 하면 도둑이 최대 금액을 훔칠 수 있을까요? 다행히 돈이 음수인 집은 없으니 최대한 많은 집을 털어야할테고 당연히 모든 집을 다 터는 게 유리할텐데요. 하지만 보안 시스템 때문에 인접하고 있는 두 집을 털면 경찰에 걸리니 조심해야겠죠?
첫 번째 집을 털면 인접하고 두 번째 집과 마지막 집을 못 털겠죠?
O X _______ X
[1, 2, 3, ..., 1]
반대로 첫 번째 집을 안 털었다면 두 번째 집부터는 자유롭게 털 수 있다는 얘기가 됩니다.
X ____________
[1, 2, 3, ..., 1]
이렇게 첫 번째 집을 털지 말지 결정하면, 다음으로 털 수 있는 첫 번째 집과 미자막 집을 털 수 있는지를 알 수 있는데요. 첫 번째 집을 털었을 경우 도둑이 훔칠 수 있는 최대 금액은, 세 번째 집부터 마지막 전 집까지 털어서 훔친 최대 금액에 첫 번째집에서 훔친 금액을 더한 것일 겁니다. 반대로 첫 번째 집을 안 털은 경우 도둑이 훔칠 수 있는 최대 금액은, 두 번째 집부터 마지막 집까지 털어서 훔진 최대 금액일 것입니다. 그리고 이 두 개의 금액 중에서 더 큰 금액이 문제에서 최종적으로 요구하는 모든 경우를 커버하는 최대 금액일 것입니다.
F()
를 집에 모여있는 돈의 배열 nums
가 주어졌을 때 훔친 수 있는 최대 금액을 구하는 함수라고 가정하고 위 로직은 다음과 같이 수식으로 나타낼 수 있습니다.
인덱스 0에 있는 집을 털었을 때 최대 금액 = nums[0] + F(nums[2:-1])
인덱스 0에 있는 집을 안 털었을 때 최대 금액 = F(nums[1:])
이렇게 크게 두 가지 경우로 분리하고 나면 House Robber 문제를 푼 것처럼 재귀 알고리즘을 적용할 수 있습니다. 입력 배열 내의 특정 인덱스부터 시작해서 훔칠 수 있는 최대 금액은 아래 점화식으로 구할 수 있습니다.
F(start) = MAX(nums[start] + F(start + 2), F(start + 1))
구현하기 까다로울 수 있는 부분이 마지막 집을 털 때 첫 번째 집을 털었는지 여부를 기억해야하는 점인데요. 이 재귀 함수의 두 번째 인자로 계속해서 호출할 때 마다 전달을 하면 될 것 같습니다.
그럼 지금까지 설명드린 재귀 알고리즘을 코드로 작성해보겠습니다.
from typing import List
class Solution:
def rob(self, nums: List[int]) -> int:
def dfs(idx, first=False):
# 더 이상 털 집이 없음
if idx >= len(nums):
return 0
# 첫 번째 집을 털었을 경우 마지막 집은 못 텀
if idx == len(nums) - 1:
return 0 if first else nums[idx]
# 인덱스 0에서 첫 번째 집을 털었을 때와 안 털었을 때 분리
if idx == 0:
return max(nums[0] + dfs(2, True), dfs(1, False))
return max(nums[idx] + dfs(idx + 2, first), dfs(idx + 1, first))
return dfs(0)
이 풀이는 매번 재귀 함수가 호출될 때 마다 두 번의 추가적인 재귀 호출이 일어나게 되므로 시간 복잡도는 O(2^n)
입니다.
반면에 주어진 배열가 길이가 길어지면 길어질 수록 재귀 호출 스택의 깊이도 그에 비례하여 깊어지므로 공간 복잡도는 O(n)
입니다.
아쉽게도 이 코드를 LeetCode에 제출하면 시간 한도를 초과하여 Time Limit Exceeded
오류가 발생하는데요.
메모이제이션(memoization) 기법을 이용해서 함수 호출의 결과를 저장해두고 재사용하여 불필요하게 반되는 연산을 제거해볼까요?
from functools import cache
class Solution:
def rob(self, nums: List[int]) -> int:
@cache
def dfs(idx, first=False):
if idx == len(nums) - 1:
return 0 if first else nums[idx]
if idx >= len(nums):
return 0
if idx == 0:
return max(nums[0] + dfs(2, True), dfs(1, False))
return max(nums[idx] + dfs(idx + 2, first), dfs(idx + 1, first))
return dfs(0)
이렇게 알고리즘을 최적화해주면 시간 복잡도를 기존 O(2^n)
에서 O(n)
로 대폭 향상시킬 수 있습니다.
이 코드를 LeetCode에 다시 제출해보면 이번에는 시간 제한이 걸리지 않고 통과할 것입니다 🤗
풀이 2: DP 1
위 재귀 풀이에서는 더 큰 배열에 대한 결과를 구하기 위해서 더 작은 배열에 대한 결과가 계산되기를 기다리는 방식으로 문제를 해결을 하는데요. 이렇게 Top-down 방향으로 문제에 접근을 해봤으니 이번에는 정반대인 Bottom-up 방향으로도 접근해볼까요?
이 문제는 동적 계획법(Dynamic Programming)이라는 테크닉을 사용하면 반복 알고리즘으로도 해결할 수 있는데요.
기본적인 아이디어는 빈 배열부터 하나 씩 배열의 크기를 늘려가면서 최대 금액을 배열에 저장해두고 작은 크기의 배열을 상대로 저장해둔 계산 결과를 더 큰 크기의 배열을 상대로 계산할 때 활용하는 것입니다.
각 배열의 크기에 대한 계산 결과는 첫 번째 집을 털었을 때는 DP1
, 첫 번째 집을 털지 않았을 때는 DP2
라는 배열에 저장해놓고 재활용하면 될 것 같아요.
그럼 문제의 두 번째 예제에서 주어진 배열인 [1, 2, 3, 1]
을 통해서 이 아이디어를 적용해보겠습니다.
먼저 털 집이 하나도 없을 때 훔칠 수 있는 최대 금액은 얼마일까요?
당연히 0
이 겠죠?
[]
DP1: [0]
DP2: [0]
배열의 크기를 한 칸 늘리면 훔칠 수 있는 최대 금액이 얼마가 될까요?
첫 번째 집을 터는 DP1
에는 해당 금액을 저장하는 반면에 첫 번째 집을 털지 않는 DP2
에는 0
을 저장합니다.
[1]
DP1: [0, 1]
DP2: [0, 0]
배열의 크기를 한 칸 더 늘려보면 어떻게 될까요?
이제 부터는 기존에 DP
배열에 구해놓은 더 작은 범위의 최대 금액을 현재 범위의 최대 금액을 구하는데 활용할 수 있습니다.
[1, 2]
MAX(DP[0] + 2, DP[1]) = MAX(0 + 2, 1) = 2
DP1: [0, 1, 2]
MAX(DP[0] + 2, DP[1]) = MAX(0 + 2, 0) = 2
DP2: [0, 0, 2]
배열의 크기를 한 칸 더 늘리면 털 수 있는 집의 범위에 세 번째 집까지 들어오게 되죠?
[1, 2, 3]
MAX(DP1[1] + 3, DP1[2]) = MAX(1 + 3, 2) = 4
DP1: [0, 1, 2, 4]
MAX(DP2[1] + 3, DP2[2]) = MAX(0 + 3, 2) = 3
DP2: [0, 0, 2, 3]
배열의 크기를 한 칸 더 늘리면 털 수 있는 집의 범위에 마지막 집까지 들어오게 되죠?
[1, 2, 3, 1]
MAX(DP1[2] + 1, DP1[3]) = MAX(2 + 1, 4) = 4
DP1: [0, 1, 2, 4, 4]
MAX(DP2[2] + 1, DP2[3]) = MAX(2 + 1, 3) = 3
DP2: [0, 0, 2, 3, 3]
최종 결과는 DP1
배열의 마지막 두 번째 값과 DP2
배열의 마지막 값을 비교해야하는데요.
DP1
배열은 첫 번째 집을 털은 경우이기 때문에 절대 마지막 집은 털 수가 없기 때문입니다.
MAX(DP1[-2], DP[-1]) = MAX(4, 3) = 4
이 DP 알고리즘을 구현해보면 다음과 같습니다.
from typing import List
class Solution:
def rob(self, nums: List[int]) -> int:
if len(nums) == 1:
return nums[0]
dp1 = [0] + [0] * len(nums)
dp1[1] = nums[0] # 첫 번째 집 텀
dp2 = [0] + [0] * len(nums)
dp2[1] = 0 # 첫 번째 집 안 텀
for n in range(2, len(dp1)):
dp1[n] = max(dp1[n - 2] + nums[n - 1], dp1[n - 1])
dp2[n] = max(dp2[n - 2] + nums[n - 1], dp2[n - 1])
return max(dp1[-2], dp2[-1])
이 풀이의 시간 복잡도와 공간 복잡도는 모두 O(n)
으로써 메모이제이션을 적용한 재귀 알고리즘과 크게 다르지는 않습니다.
풀이 3: DP 2
위 동적 계획법 알고리즘은 입력 배열의 길이와 동일한 2개의 배열을 사용하는데요. 훔칠 수 있는 최대 금액을 굳이 모든 크기의 배열에 대해서 다 저장해 둘 필요가 있을까요?
사실 배열의 크기가 1일 때 계산 결과와 배열의 크기가 2일 때 계산 결과는 크기가 3인 배열에 대한 결과를 구할 때 활용하면 그 이후로는 쓸모가 없습니다.
따라서 우리는 현재 배열의 크기가 보다 1과 2가 작은 크기의 배열에 대한 계산 결과만 저장해두면 된다는 것을 깨닫게 되며,
즉, 길이가 n + 1
인 배열 대신에 단 2개의 변수만으로도 이 알고리즘은 구현할 수 있을 것입니다.
하지만 2개의 배열이 없이 모든 집들은 원형으로 배치되어 있는 부분은 어떻게 처리해야할까요? 여기서 한 가지 아이디어를 내보면, 첫 번째 집을 털면 절대 마지막 집을 털 가능성이 없습니다. 따라서 첫 번째 집을 터는 경우에는 아예 마지막 집을 제외하고 최대 금액을 구합니다.
___________________
[1, 2, 3, 1, 2, 3, 1, 2]
마찬가지로 마지막 집을 털면 첫 번째 집을 털 가능성이 없죠. 따라서 마지막 집을 터는 경우에는 아예 첫 번째 집을 제외하고 최대 금액을 구합니다.
___________________
[1, 2, 3, 1, 2, 3, 1, 2]
이 두 개의 최대 금액 중 더 큰 금액이 우리가 구하고자 하는 최대 금액이 될 것입니다.
MAX(dp(nums[1:]), dp(nums[:-1]))
그럼 이 알고리즘을 코드로 구현해볼까요?
파이썬에서 리스트를 슬라이싱(slicing)하는데 O(n)
의 메모리가 소모되므로 재귀 함수가 시작 인덱스와 종료 인덱스를 인자로 받도록 한 부분 주의깊게 보시길 바랄께요.
class Solution:
def rob(self, nums: List[int]) -> int:
if len(nums) == 1:
return nums[0]
def dp(start, end):
prev, curr = 0, 0
for idx in range(start, end):
prev, curr = curr, max(prev + nums[idx], curr)
return curr
return max(dp(0, len(nums) - 1), dp(1, len(nums)))
이렇게 알고리즘을 최적화해주면 공간 복잡도를 O(1)
로 내릴 수 있습니다.
이전에 작성한 코드도 나쁘지는 않지만 이렇게 추가적으로 최적화까지 해주면 코딩 테스트에서 더 좋은 결과를 얻을 수 있을 것입니다.
마치면서
마지막 풀이에서 사용한 아이디어는 직관(intuition)이 필요하며 코딩 문제를 많이 푸시다보시면 자연스럽게 길러지게 됩니다. 그러니 아이디어를 직접 떠올리지 못하셨더라도 좌절하시지 않으셨으면 좋겠습니다.
이 문제가 너무 어려우셨다면 비슷하지만 좀 더 쉬운 문제인 House Robber를 풀어보시라고 추천드립니다.