파이썬으로 인접한 나무의 높이가 모두 다르도록 만드는 최소 비용 구하기

문제 소개

식물의 높이를 담은 숫자 리스트 heights와, 각 식물의 높이를 1만큼 올릴 때 드는 비용을 담은 리스트 costs가 주어집니다. 이때 인접한 식물끼리는 서로 다른 높이를 갖도록 만들어야 하며, 그렇게 하기 위해 필요한 최소 비용을 구하는 것이 목표입니다.

예를 들어 heights = [3, 2, 2], costs = [2, 5, 3]라고 한다면 정답은 3입니다. 세 번째 나무의 높이를 1만큼 올리면 비용이 3이 들고, 그 결과 높이는 [3, 2, 3]이 되어 인접한 나무들의 높이가 모두 달라지기 때문입니다.

풀이 접근 방법

이 문제는 재귀 호출을 기반으로 한 동적 계획법(DP)으로 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

각 나무의 높이를 0, 1, 2만큼 올리는 세 가지 경우만 시도해 봅니다. 직전 나무의 조정된 높이와 비교했을 때 이 세 가지 중 반드시 하나 이상은 높이가 겹치지 않으므로 충분합니다.
현재 나무의 새로운 높이가 바로 앞 나무의 높이와 다른 경우에만 다음 나무로 진행합니다.
마지막 나무에 도달하면 직전 높이와 같은지만 확인하고, 같다면 해당 나무를 1만큼 올리는 비용을 결과에 더합니다.

알고리즘 단계

dp(idx, l_height) 함수를 정의합니다. 여기서 idx는 현재 나무의 인덱스, l_height는 바로 앞 나무의 조정된 높이입니다.
idx가 마지막 인덱스(len(heights) - 1)라면, 현재 높이가 l_height와 다를 때는 0을, 같을 때는 costs[idx]를 반환합니다.
그 외의 경우에는 ret을 무한대(inf)로 초기화합니다.
i를 0부터 2까지 순회하면서, heights[idx] + i가 l_height와 다르면 ret을 dp(idx + 1, heights[idx] + i) + costs[idx] * i 값과 비교해 더 작은 값으로 갱신합니다.
ret을 반환하고, 메인에서는 dp(0, None)을 호출하여 최종 결과를 얻습니다.

구현 코드

class Solution:
    def solve(self, heights, costs):
        def dp(idx, l_height):
            if idx == len(heights) - 1:
                return 0 if heights[idx] != l_height else costs[idx]
            ret = float("inf")
            for i in range(3):
                if heights[idx] + i != l_height:
                    ret = min(ret, dp(idx + 1, heights[idx] + i) + costs[idx] * i)
            return ret
        return dp(0, None)

ob = Solution()
heights = [3, 2, 2]
costs = [2, 5, 3]
print(ob.solve(heights, costs))

실행 결과

입력:

[3, 2, 2], [2, 5, 3]

출력:

성능 개선 팁

위 재귀 구현은 입력 크기가 커지면 지수적으로 많은 호출이 발생할 수 있습니다. dp 함수에 @functools.lru_cache 데코레이터를 붙이거나 딕셔너리로 메모이제이션을 적용하면 각 상태를 한 번씩만 계산하여 O(n × 3) 시간 안에 효율적으로 처리할 수 있습니다.