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
# 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]