LeetCode/DP심화

309. Best Time to Buy and Sell Stock with Cooldown ★★★

hyunkookim 2025. 1. 7. 19:08
from typing import List

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        # 최대 이익을 계산해서, return
        # 언제든지 사고팔수 있지만,
        # 살려면. 기존 주식을 팔아야 하고
        # 팔면 그다음날 cooldown 후, 그 다음날 살수 있음

        """
        root -> buy -> sell -> cooldown -> buy      -> sell
                    -> cooldown         |           -> cooldown
             -> cooldown -> buy         -> cooldown
                         -> cooldown
        """
        # 이렇게 하면, 트리의 높이는 n 이 됨
        # O(n)

        # State: Buying (true) or Selling (false)
        # if Buy -> i+1
        # if Sell -> i+2 왜냐면, 사고나서 cooldown day 후에 살수있어서, +2

        dp = {} # 캐슁->hashmap, key=(i, buying) val= max_profit

        def dfs(i, buying):
            if i >= len(prices):  # 현재 날짜가 prices 길이보다 크거나 같으면 종료
                return 0  # 더 이상 행동할 수 없으므로 이익은 0 반환

            if (i, buying) in dp:  # 이미 계산된 값이 메모이제이션에 있다면 반환
                return dp[(i, buying)]

            cooldown = dfs(i + 1, buying)  # 현재 상태를 유지하며 다음 날로 넘어감 (cooldown)

            if buying:  # 주식을 구매하는 상태일 때
                buy = dfs(i + 1, not buying) - prices[i]  # 주식을 구매하고 다음 날로 이동
                dp[(i, buying)] = max(buy, cooldown)  # 구매하거나 쿨다운 중 더 큰 값을 저장
            else:  # 주식을 판매하는 상태일 때 Selling
                sell = dfs(i + 2, not buying) + prices[i]  # 주식을 판매하고 2일 후로 이동, 이때 not buying 는 당연히 true
                dp[(i, buying)] = max(sell, cooldown)  # 판매하거나 쿨다운 중 더 큰 값을 저장

            return dp[(i, buying)]  # 현재 상태에서 최대 이익 반환

        return dfs(0, True)  # 첫 번째 날부터 시작하며, 주식을 구매하는 상태로 시작