LeetCode 309: Best Time To Buy And Sell Stock With Cooldown — Step-by-Step Visual Trace

#leetcode #algorithms #python #dynamicprogramming

Medium — Dynamic Programming | Array | State Machine

The Problem

Find the maximum profit from buying and selling stocks with a cooldown period of one day after each sale. You can complete multiple transactions but must wait one day after selling before buying again.

Approach

Use dynamic programming with three states: buy (maximum profit after buying), sell (maximum profit after selling), and cooldown (maximum profit during cooldown). Track the optimal profit for each state at every day by considering all possible transitions.

Time: O(n) · Space: O(1)

Code

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        if not prices:
            return 0

        # Initialize variables to represent the maximum profit after each action.
        buy = -prices[
            0
        ]  # Maximum profit after buying on day 0 (negative because we spend money).
        sell = 0  # Maximum profit after selling on day 0 (no profit yet).
        cooldown = 0  # Maximum profit after cooldown on day 0 (no profit yet).

        for i in range(1, len(prices)):
            # To maximize profit on day 'i', we can either:

            # 1. Buy on day 'i'. We subtract the price of the stock from the maximum profit after cooldown on day 'i-2'.
            new_buy = max(buy, cooldown - prices[i])

            # 2. Sell on day 'i'. We add the price of the stock to the maximum profit after buying on day 'i-1'.
            new_sell = buy + prices[i]

            # 3. Do nothing (cooldown) on day 'i'. We take the maximum of the maximum profit after cooldown on day 'i-1' and after selling on day 'i-1'.
            new_cooldown = max(cooldown, sell)

            # Update the variables for the next iteration.
            buy, sell, cooldown = new_buy, new_sell, new_cooldown

        # The maximum profit will be the maximum of the profit after selling on the last day and the profit after cooldown on the last day.
        return max(sell, cooldown)