Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 최대 k번의 매수·매도에서 얻을 수 있는 최대 수익 구하기

시간 순서대로 정렬된 주가 목록 nums와 정수 k가 주어졌을 때, 최대 k번의 매수와 매도를 통해 얻을 수 있는 최대 수익을 구하는 프로그램을 만들어 보겠습니다. 단, 반드시 매수 후에 매도해야 하며, 매도한 뒤에만 다시 매수할 수 있다는 조건이 있습니다.

예를 들어 prices = [7, 3, 5, 2, 3]이고 k = 2라면 출력은 3이 됩니다. 가격이 3일 때 매수하여 5일 때 매도하고(+2), 다시 2일 때 매수하여 3일 때 매도하면(+1) 총 수익이 3이 되기 때문입니다.

문제 해결 접근 방식

이 문제는 재귀 호출과 동적 계획법(DP)을 활용해 해결할 수 있습니다. 핵심 아이디어는 각 날짜를 기준으로 '주식을 보유한 상태'와 '보유하지 않은 상태'로 나누어, 각 상태에서 가능한 선택지 중 더 큰 이익을 재귀적으로 계산하는 것입니다.

  • dp(i, k, bought) 함수를 정의합니다. 여기서 i는 현재 날짜 인덱스, k는 남은 거래 횟수, bought는 현재 주식을 보유 중인지 여부를 나타냅니다.
  • i가 prices의 길이와 같거나 k가 0이면 더 이상 거래할 수 없으므로 0을 반환합니다.
  • bought가 참인 경우(주식 보유 중): max(dp(i+1, k-1, False) + prices[i], dp(i+1, k, bought))를 반환합니다. 즉, 오늘 매도할지 아니면 계속 보유할지 중 더 큰 값을 선택합니다.
  • bought가 거짓인 경우(주식 미보유): max(dp(i+1, k, True) - prices[i], dp(i+1, k, bought))를 반환합니다. 즉, 오늘 매수할지 아니면 관망할지 중 더 큰 값을 선택합니다.
  • 메인 메서드에서 dp(0, k, False)를 호출하고 그 결과를 반환합니다.

아래 예제를 통해 더 자세히 살펴보겠습니다.

구현 예제

class Solution:
    def solve(self, prices, k):
        def dp(i, k, bought):
            if i == len(prices) or k == 0:
                return 0
            if bought:
                return max(dp(i + 1, k - 1, False) + prices[i], dp(i + 1, k, bought))
            else:
                return max(dp(i + 1, k, True) - prices[i], dp(i + 1, k, bought))

        return dp(0, k, False)

ob = Solution()
prices = [7, 3, 5, 2, 3]
k = 2
print(ob.solve(prices, k))

입력

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

출력

3

복잡도 분석 및 성능 개선 팁

위 재귀 구현의 상태 공간은 O(n × k × 2)이므로, 메모이제이션(memoization)을 적용하면 시간 복잡도 O(n × k), 공간 복잡도 O(n × k)로 개선할 수 있습니다. 파이썬에서는 functools.lru_cache 데코레이터를 dp 함수에 붙여주는 것만으로 간단하게 중복 계산을 제거할 수 있으며, 입력 크기가 커질 때 실행 속도가 크게 향상됩니다.