파이썬으로 주식 매매 최대 이익 구하기: 쿨다운 기간이 있는 주식 거래 문제 풀이

문제 설명

어떤 회사의 주식 가격이 시간 순서대로 리스트로 주어졌다고 가정해 봅시다. 이때 주식을 사고팔아 얻을 수 있는 최대 이익을 구해야 합니다. 단, 다음 두 가지 제약 조건이 있습니다.

• 주식을 팔기 전에 반드시 먼저 매수해야 합니다.

• 주식을 매도한 후에는 하루를 기다린 뒤에야 다시 매수할 수 있습니다(쿨다운 규칙).

예를 들어 입력이 prices = [2, 6, 9, 4, 11]이라면 출력은 11이 됩니다. 2원에 사서 6원에 팔고, 하루를 기다린 뒤 4원에 다시 사서 11원에 팔면 되기 때문입니다.

풀이 접근 방법

이 문제는 동적 계획법(Dynamic Programming)으로 해결할 수 있습니다. 핵심은 각 날짜마다 두 가지 상태를 추적하는 것입니다.

• s: 주식을 보유하고 있지 않은(매도 완료) 상태에서의 최대 이익 → 초기값 0

• b: 주식을 보유하고 있는 상태에서의 최대 이익 → 초기값 음의 무한대(-∞)

각 날짜 i에 대해 다음 절차를 수행합니다.

• temp := b — 현재 보유 상태 값을 임시 저장합니다.

• b := max(b, s - prices[i]) — 오늘 새로 매수하는 경우와 기존 보유 상태 중 더 나은 쪽을 선택합니다.

• i가 0이 아니라면 s := max(s, temp + prices[i - 1]) — 어제 매도했을 때의 이익을 반영합니다. 이 과정 덕분에 오늘 매수한 주식을 어제 판 것으로 처리하지 않아 하루 쿨다운 규칙이 자연스럽게 적용됩니다.

마지막으로 s와 b + 마지막 날 가격 중 더 큰 값을 반환하면 정답이 됩니다.

구현 예제

class Solution:
    def solve(self, prices):
        s = 0
        b = float("-inf")
        for i in range(len(prices)):
            temp = b
            b = max(b, s - prices[i])
            if i:
                s = max(s, temp + prices[i - 1])
        return max(s, b + prices[-1])

ob = Solution()
prices = [2, 6, 9, 4, 11]
print(ob.solve(prices))

입력

[2, 6, 9, 4, 11]

출력

11

복잡도 분석

• 시간 복잡도: O(n) — 가격 리스트를 한 번만 순회합니다.

• 공간 복잡도: O(1) — 두 개의 변수(s, b)만 사용합니다.