Best Time to Buy and Sell Stock with Cooldown

LeetCode 309: Best Time to Buy and Sell Stock with Cooldown Solution & Explanation

Problem Overview

TL;DR: Modeled as a state machine, the optimal solution tracks the maximum profit across three distinct states (holding, sold/cooldown, and reset) for each day to handle the transaction restrictions efficiently.

The Best Time to Buy and Sell Stock with Cooldown problem asks us to calculate the maximum profit achievable from a series of stock prices. We can buy and sell multiple times, but there is a strict constraint: after selling a stock, we must wait one day (cooldown) before buying again. Additionally, we cannot engage in multiple transactions simultaneously; we must sell our current holding before buying a new one.

This is a classic variation of the "Best Time to Buy and Sell Stock" series found in LeetCode 309, requiring a strategic approach to handle the inter-day dependency introduced by the cooldown period.

Brute Force Approach for Best Time to Buy and Sell Stock with Cooldown

The brute force approach involves exploring every possible sequence of actions (Buy, Sell, Cooldown) for every day. This can be visualized as a decision tree where, at each step, we branch out based on valid moves.

1. If we are holding a stock, we can either Sell or Wait.
2. If we are not holding a stock and not in cooldown, we can Buy or Wait.
3. If we are in cooldown, we must Wait (forcing a state reset).

We would implement this using recursion to calculate the profit for every path and return the maximum.

1# Pseudo-code for Brute Force
2def solve(day, holding, cooldown):
3    if day >= len(prices):
4        return 0
5    
6    # Option 1: Do nothing (skip day)
7    profit = solve(day + 1, holding, False)
8    
9    if holding:
10        # Option 2: Sell (triggers cooldown next)
11        profit = max(profit, prices[day] + solve(day + 1, False, True))
12    elif not cooldown:
13        # Option 3: Buy (only if not holding and not in cooldown)
14        profit = max(profit, -prices[day] + solve(day + 1, True, False))
15        
16    return profit

Time Complexity: $O(2^N)$, where $N$ is the number of days. In the worst case, the decision tree branches exponentially. Why it fails: With prices.length up to 5000, an exponential solution will inevitably result in a Time Limit Exceeded (TLE) error. The recursion repeatedly calculates the same subproblems (e.g., "maximum profit from day 10 given we are currently holding stock").

Algorithm Strategy: Dynamic Programming (DP) Patterns

We utilize the identified pattern by defining three variables (or arrays) to represent the maximum profit ending in each state for the current day.

1. State Definitions:

   • held[i]: Max profit on day i ending with stock in hand.
   • sold[i]: Max profit on day i ending immediately after selling stock.
   • reset[i]: Max profit on day i ending with no stock and no immediate sale (ready to buy).

2. State Transitions:

   • To be in the held state today, we either:
     • Kept holding from yesterday (held[i-1]).
     • Bought stock today. To buy, we must have been in the reset state yesterday (reset[i-1] - price). Note: We cannot buy from the sold state due to cooldown.
   • To be in the sold state today, we must have:
     • Sold the stock we were holding yesterday (held[i-1] + price).
   • To be in the reset state today, we either:
     • Continued resting from yesterday (reset[i-1]).
     • Just finished a cooldown (meaning we were in the sold state yesterday) (sold[i-1]).

3. Initialization:

   • held: Starts at -prices[0] (we bought stock).
   • sold: Starts at 0 (or theoretically impossible, often initialized to minimum integer, but 0 works if logic is careful).
   • reset: Starts at 0.

4. Result:

   • The answer is the maximum of sold and reset on the last day. We do not consider held because ending with stock in hand is never more profitable than selling it or not buying it in the last step.

Execution Flow

1. Initialize Variables:

   • held = -prices[0]
   • sold = 0
   • reset = 0

2. Iterate from the second day (index 1) to the end of the array:

   • Store the previous values of held, sold, and reset to avoid using updated values from the current day in the same calculation step.
   • Update held: max(previous_held, previous_reset - current_price)
     • Logic: Either keep holding, or buy new stock (only allowed from reset state).
   • Update sold: previous_held + current_price
     • Logic: We sell the stock we were holding.
   • Update reset: max(previous_reset, previous_sold)
     • Logic: Either keep waiting, or recover from the cooldown of a previous sale.

3. Termination:

   • After the loop finishes, return max(sold, reset).

Proof of Correctness

The algorithm is correct by induction.

• Base Case: On day 0, the states are correctly initialized (buying costs money, resting costs nothing).
• Inductive Step: Assume that on day i-1, the variables held, sold, and reset correctly hold the maximum possible profit for those specific states.
  • The transition to held[i] considers all valid ways to have stock (wait or buy).
  • The transition to sold[i] considers the only valid way to sell (must have stock).
  • The transition to reset[i] considers all valid ways to be empty-handed without selling today (wait or recover from cooldown).
• Since the transitions cover the entire state space and enforce the cooldown constraint (buying only from reset, not sold), the values at day N must be optimal.

Pattern Reuse Notes

The DP - Stock problems pattern is versatile. The core logic of state transitions applies to several related LeetCode problems:

• LeetCode 121: Best Time to Buy and Sell Stock: The base case with $k=1$ transaction.
• LeetCode 122: Best Time to Buy and Sell Stock II: Infinite transactions without cooldown.
• LeetCode 123: Best Time to Buy and Sell Stock III: Exactly $k=2$ transactions.
• LeetCode 188: Best Time to Buy and Sell Stock IV: Generalized to $k$ transactions.

The primary difference in LeetCode 309 is strictly the modification of the state transition diagram to include the sold -> reset delay.

Reference Implementation

C++ Solution for LeetCode 309

1#include <vector>
2#include <algorithm>
3#include <climits>
4
5using namespace std;
6
7class Solution {
8public:
9    int maxProfit(vector<int>& prices) {
10        if (prices.empty()) return 0;
11
12        // held: max profit if we end the day holding a stock
13        // sold: max profit if we end the day having just sold a stock
14        // reset: max profit if we end the day not holding and not just sold
15        
16        int held = -prices[0];
17        int sold = 0;
18        int reset = 0;
19
20        for (int i = 1; i < prices.size(); ++i) {
21            int prev_held = held;
22            int prev_sold = sold;
23            int prev_reset = reset;
24
25            // Transition logic
26            held = max(prev_held, prev_reset - prices[i]);
27            sold = prev_held + prices[i];
28            reset = max(prev_reset, prev_sold);
29        }
30
31        return max(sold, reset);
32    }
33};

Java Solution for LeetCode 309

1class Solution {
2    public int maxProfit(int[] prices) {
3        if (prices == null || prices.length == 0) {
4            return 0;
5        }
6
7        // Initialize states for day 0
8        int held = -prices[0];
9        int sold = 0;
10        int reset = 0;
11
12        for (int i = 1; i < prices.length; i++) {
13            int prevHeld = held;
14            int prevSold = sold;
15            int prevReset = reset;
16
17            // Update states based on previous day
18            held = Math.max(prevHeld, prevReset - prices[i]);
19            sold = prevHeld + prices[i];
20            reset = Math.max(prevReset, prevSold);
21        }
22
23        // The result is the max of selling today or having rested
24        return Math.max(sold, reset);
25    }
26}

Python Solution for LeetCode 309

1class Solution:
2    def maxProfit(self, prices: list[int]) -> int:
3        if not prices:
4            return 0
5        
6        # State initialization
7        held = -prices[0]
8        sold = 0
9        reset = 0
10        
11        for i in range(1, len(prices)):
12            prev_held = held
13            prev_sold = sold
14            prev_reset = reset
15            
16            # State transitions
17            # 1. Hold: Keep holding OR Buy (only possible from reset)
18            held = max(prev_held, prev_reset - prices[i])
19            
20            # 2. Sold: Must sell from held state
21            sold = prev_held + prices[i]
22            
23            # 3. Reset: Keep resting OR Cooldown over (from sold)
24            reset = max(prev_reset, prev_sold)
25            
26        return max(sold, reset)