Best Time to Buy and Sell Stock with Cooldown

MediumDynamic ProgrammingState Machine
Category: 2-D Dynamic Programming
Companies that ask this question:
AmazonGoogleFacebookMicrosoftBloomberg

Approach

Best Time to Buy and Sell Stock with Cooldown

Problem Statement

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).

Examples

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

Approach: State Machine Dynamic Programming

Intuition

This problem is best solved using a state machine approach. At any point, we can be in one of three states:

  1. Hold: We are holding a stock
  2. Sold: We just sold a stock (must cooldown next day)
  3. Rest: We are not holding and can buy

The key insight is the cooldown constraint: after selling, we must rest for one day before buying again.

State Transitions

hold → sold  (sell stock)
sold → rest  (mandatory cooldown)
rest → hold  (buy stock)
rest → rest  (continue resting)
hold → hold  (keep holding)

Algorithm

Define three arrays:

  • hold[i] = max profit on day i if holding stock
  • sold[i] = max profit on day i if just sold
  • rest[i] = max profit on day i if resting (not holding)

Recurrence Relations:

hold[i] = max(hold[i-1], rest[i-1] - prices[i])
sold[i] = hold[i-1] + prices[i]
rest[i] = max(rest[i-1], sold[i-1])

Base Case:

hold[0] = -prices[0]  # Buy on day 0
sold[0] = 0           # Can't sell on day 0
rest[0] = 0           # Start with no profit

Final Answer:

max(sold[n-1], rest[n-1])  # Can't end holding stock

Why This Works

  • hold[i]: We either kept holding from yesterday, or bought today from rest state
  • sold[i]: We sell the stock we were holding
  • rest[i]: We either kept resting, or completed cooldown after selling yesterday
  • The cooldown is automatically enforced because we can only buy from the rest state

Complexity Analysis

  • Time Complexity: O(n) - single pass through prices
  • Space Complexity: O(n) for arrays, can be optimized to O(1)

Space Optimization

We only need the previous day's values, so we can use just 6 variables instead of 3 arrays:

hold, sold, rest = -prices[0], 0, 0

for price in prices[1:]:
    prev_hold, prev_sold, prev_rest = hold, sold, rest
    hold = max(prev_hold, prev_rest - price)
    sold = prev_hold + price
    rest = max(prev_rest, prev_sold)

return max(sold, rest)

This reduces space to O(1).

Common Patterns

This problem demonstrates:

  1. State Machine DP: Multiple states with defined transitions
  2. Constraint Handling: Cooldown enforced through state design
  3. Space Optimization: Reducing from O(n) to O(1)

Related Problems

  • Best Time to Buy and Sell Stock (basic version)
  • Best Time to Buy and Sell Stock II (unlimited transactions)
  • Best Time to Buy and Sell Stock III (at most 2 transactions)
  • Best Time to Buy and Sell Stock with Transaction Fee

Tips

  1. Draw the state machine diagram to visualize transitions
  2. Remember: can't end in "hold" state (no profit if still holding)
  3. The "sold" state automatically creates the cooldown requirement
  4. Space optimization: only need previous values, not entire arrays

Solution

java
1class Solution {
2    // Solution 1: State Machine DP (3 arrays)
3    public int maxProfit(int[] prices) {
4        if (prices == null || prices.length == 0) {
5            return 0;
6        }
7        
8        int n = prices.length;
9        int[] hold = new int[n];  // Holding stock
10        int[] sold = new int[n];  // Just sold stock
11        int[] rest = new int[n];  // Not holding, can buy
12        
13        // Initial state
14        hold[0] = -prices[0];
15        sold[0] = 0;
16        rest[0] = 0;
17        
18        for (int i = 1; i < n; i++) {
19            // Hold: either keep holding or buy from rest
20            hold[i] = Math.max(hold[i-1], rest[i-1] - prices[i]);
21            // Sold: must come from hold
22            sold[i] = hold[i-1] + prices[i];
23            // Rest: either continue resting or from sold (cooldown)
24            rest[i] = Math.max(rest[i-1], sold[i-1]);
25        }
26        
27        // Can't end in hold state
28        return Math.max(sold[n-1], rest[n-1]);
29    }
30    
31    // Solution 2: Space-Optimized DP (O(1) space)
32    public int maxProfitOptimized(int[] prices) {
33        if (prices == null || prices.length == 0) {
34            return 0;
35        }
36        
37        int hold = -prices[0];
38        int sold = 0;
39        int rest = 0;
40        
41        for (int i = 1; i < prices.length; i++) {
42            int prevHold = hold;
43            int prevSold = sold;
44            int prevRest = rest;
45            
46            hold = Math.max(prevHold, prevRest - prices[i]);
47            sold = prevHold + prices[i];
48            rest = Math.max(prevRest, prevSold);
49        }
50        
51        return Math.max(sold, rest);
52    }
53    
54    // Solution 3: Alternative DP formulation
55    public int maxProfitAlternative(int[] prices) {
56        if (prices == null || prices.length == 0) {
57            return 0;
58        }
59        
60        int buy = -prices[0];
61        int sell = 0;
62        int cooldown = 0;
63        
64        for (int i = 1; i < prices.length; i++) {
65            int prevBuy = buy;
66            int prevSell = sell;
67            int prevCooldown = cooldown;
68            
69            buy = Math.max(prevBuy, prevCooldown - prices[i]);
70            sell = prevBuy + prices[i];
71            cooldown = Math.max(prevCooldown, prevSell);
72        }
73        
74        return Math.max(sell, cooldown);
75    }
76    
77    // Solution 4: Using buy/sell with 2-day lookback
78    public int maxProfitTwoDay(int[] prices) {
79        if (prices == null || prices.length == 0) {
80            return 0;
81        }
82        
83        int n = prices.length;
84        int[] buy = new int[n];
85        int[] sell = new int[n];
86        
87        buy[0] = -prices[0];
88        sell[0] = 0;
89        
90        for (int i = 1; i < n; i++) {
91            // Can buy if we didn't sell yesterday
92            if (i >= 2) {
93                buy[i] = Math.max(buy[i-1], sell[i-2] - prices[i]);
94            } else {
95                buy[i] = Math.max(buy[i-1], -prices[i]);
96            }
97            
98            // Can always sell if we're holding
99            sell[i] = Math.max(sell[i-1], buy[i-1] + prices[i]);
100        }
101        
102        return sell[n-1];
103    }
104}
105
Loading visualizer...