Read-only web browser

Solutions

Browse the interview solutions folder with syntax highlighting.

best-time-to-buy-and-sell-stock-with-cooldown.py

2d-dp/best-time-to-buy-and-sell-stock-with-cooldown.py · Python · 2.1 KiB · 2026-05-03 15:45

Back to folder
# DP top down or bottom up. Check all cases: skip, buy, sell. 
# dp

# Top-down DP - O(n) time, O(n) space
def maxProfit(prices: list[int]) -> int:
    dp = {}

    def dfs(day: int, buyprice: int) -> int:
        if day >= len(prices):
            return 0
        if (day, buyprice) in dp:
            return dp[(day,buyprice)]

        res = 0

        # skipping the day is always an option 
        skip = dfs(day+1, buyprice)

        # we're holding something we bought
        if buyprice >= 0: 
            sell = prices[day] - buyprice + dfs(day+2, -1) # +2 days to honor cooldown
            res = max(sell, skip)
        # we don't own anything
        else:
            buy = dfs(day+1, prices[day])
            res = max(buy, skip)
        dp[(day,buyprice)] = res
        return res

    return dfs(0, -1)


# Top-down DP (bool version) - O(n) time, O(n) space
def maxProfit(prices: list[int]) -> int:
    dp = {}  # key=(i, buying) val=max_profit

    def dfs(i: int, buying: bool) -> int:
        if i >= len(prices):
            return 0
        if (i, buying) in dp:
            return dp[(i, buying)]

        cooldown = dfs(i + 1, buying)
        if buying:
            buy = dfs(i + 1, False) - prices[i]
            dp[(i, buying)] = max(buy, cooldown)
        else:
            sell = dfs(i + 2, True) + prices[i]
            dp[(i, buying)] = max(sell, cooldown)
        return dp[(i, buying)]

    return dfs(0, True)


# Bottom up DP - O(n) time, O(n) space
def maxProfit(prices: list[int]) -> int:
    n = len(prices)
    # dp[i][0 or 1]: max profit starting from day if we're allowed to buy (1) or not (0)
    dp = [[0, 0] for _ in range(n +1)]

    for i in range(n - 1, -1, -1):
        for buying in [True, False]:
            if buying:
                buy = dp[i + 1][False] - prices[i] if i + 1 < n else -prices[i]
                cooldown = dp[i + 1][True] if i + 1 < n else 0
                dp[i][1] = max(buy, cooldown)
            else:
                sell = dp[i + 2][True] + prices[i] if i + 2 < n else prices[i]
                cooldown = dp[i + 1][False] if i + 1 < n else 0
                dp[i][0] = max(sell, cooldown)

    return dp[0][1]