309. Best Time to Buy and Sell Stock with Cooldown
# O(n) time || O(1) space
def max_profit(self, prices: List[int]) -> int:
buy, sell, cooldown = -prices[0], 0, 0
for price in prices[1:]:
_buy = max(buy, cooldown - price)
_sell = buy + price
_cooldown = max(cooldown, sell)
buy, sell, cooldown = _buy, _sell, _cooldown
return max(sell, cooldown)
For this LeetCode problem, the challenge is to find the maximum profit from buying and selling stocks with a constraint: after selling a stock, you must have a cooldown period of one day before you can buy again. To solve this problem, we can use dynamic programming, where we maintain three states at each day:
- Buy: The maximum profit on day
iif we buy the stock on this day. This state can only be reached from the cooldown state of the previous day. - Sell: The maximum profit on day
iif we sell the stock on this day. This state can be reached from the buy state of the previous day. - Cooldown: The maximum profit on day
iif we are in the cooldown state on this day. This state can be reached either from the sell state or the cooldown state of the previous day.
Here’s how you can implement the solution:
-
Initialize Variables:
buy: Initially set to-prices[0]because we spendprices[0]money to buy the stock.sell: Initially set to 0, as we haven’t made any profit yet.cooldown: Initially set to 0, as we haven’t sold any stock yet.
-
Iterate Through the Array:
- For each day, calculate the new values of
buy,sell, andcooldown:new_buyis the maximum of the previousbuyand thecooldownfrom the previous day minus today’s price ( since buying costs money).new_sellis the maximum of the previousselland thebuyfrom the previous day plus today’s price (since selling earns money).new_cooldownis the maximum of the previouscooldownand thesellfrom the previous day.
- Update
buy,sell, andcooldownwith these new values for the next iteration.
- For each day, calculate the new values of
-
Return the Maximum of
sellandcooldown:- At the end, the maximum profit will be the maximum of the
sellandcooldownstates, as you can’t end with a buy.
- At the end, the maximum profit will be the maximum of the
# O(n) time || O(n) space
def max_profit_top_down(self, prices: List[int]) -> int:
@lru_cache(None)
def dp(i, holding):
if i >= len(prices):
return 0
if holding:
res = max(prices[i] + dp(i + 2, False), dp(i + 1, True))
else:
res = max(-prices[i] + dp(i + 1, True), dp(i + 1, False))
return res
return dp(0, False)