View on NeetCode
View on LeetCode

Problem

You are given an array prices where prices[i] is the price of a given stock on the ith day.

Find the maximum profit you can achieve. You may complete as many transactions as you like (i.e., buy one and sell one share of the stock multiple times) with the following restrictions:

• After you sell your stock, you cannot buy stock on the next day (i.e., cooldown one day).

Note: You may not engage in multiple transactions simultaneously (i.e., you must sell the stock before you buy again).

Example 1:

Input: prices = [1,2,3,0,2]
Output: 3
Explanation: transactions = [buy, sell, cooldown, buy, sell]

Example 2:

Input: prices = [1]
Output: 0

Constraints:

• 1 <= prices.length <= 5000
• 0 <= prices[i] <= 1000

Solution

Approach 1: State Machine DP

The key insight is to model three states: holding stock, sold (cooldown), and not holding (ready to buy).

Implementation

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

        n = len(prices)

        # Three states:
        # hold: max profit when holding stock
        # sold: max profit when just sold (must cooldown)
        # rest: max profit when not holding and can buy

        hold = -prices[0]  # Buy on day 0
        sold = 0
        rest = 0

        for i in range(1, n):
            prev_hold = hold
            prev_sold = sold
            prev_rest = rest

            # Can hold by keeping previous hold or buying from rest
            hold = max(prev_hold, prev_rest - prices[i])

            # Can be in sold state by selling from hold
            sold = prev_hold + prices[i]

            # Can be in rest by staying in rest or from cooldown after sold
            rest = max(prev_rest, prev_sold)

        # Maximum profit is either in sold or rest state
        return max(sold, rest)

Approach 2: Simplified DP

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        if len(prices) <= 1:
            return 0

        n = len(prices)

        # buy[i]: max profit on day i ending with a buy
        # sell[i]: max profit on day i ending with a sell
        buy = [0] * n
        sell = [0] * n

        buy[0] = -prices[0]
        buy[1] = max(-prices[0], -prices[1])
        sell[1] = max(0, buy[0] + prices[1])

        for i in range(2, n):
            # Buy today or keep previous buy state
            buy[i] = max(buy[i-1], sell[i-2] - prices[i])

            # Sell today or keep previous sell state
            sell[i] = max(sell[i-1], buy[i-1] + prices[i])

        return sell[n-1]

Approach 3: Space-Optimized

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

        # sold: profit after selling (must cooldown next)
        # hold: profit while holding stock
        # rest: profit while resting (can buy)

        sold = float('-inf')
        hold = float('-inf')
        rest = 0

        for price in prices:
            prev_sold = sold
            prev_hold = hold
            prev_rest = rest

            hold = max(prev_hold, prev_rest - price)
            sold = prev_hold + price
            rest = max(prev_rest, prev_sold)

        return max(sold, rest)

Complexity Analysis

• Time Complexity: O(n), where n is the number of days. Single pass through prices.
• Space Complexity: O(1) for space-optimized versions, O(n) for array-based versions.

Key Insights

1. State Machine: Three states: holding stock, just sold (cooldown required), and resting (can buy).

2. Cooldown Constraint: After selling, must wait one day before buying again. This is captured by buying from “rest” state, which comes from “sold - 1”.

3. Transitions:
   • Hold: Keep holding or buy from rest state
   • Sold: Sell from hold state
   • Rest: Stay in rest or cooldown from sold

4. No Simultaneous Transactions: Can only hold one stock at a time, enforced by state transitions.

5. Final Answer: Maximum of sold and rest states (not holding at end).

6. Cooldown Implementation: buy[i] = max(buy[i-1], sell[i-2] - prices[i]) - buy from profit two days ago (after cooldown).