261. Best Time to Buy and Sell Stock with Cooldown

🔗 LeetCode Problem: 309. Best Time to Buy and Sell Stock with Cooldown
📊 Difficulty: Medium
🏷️ Topics: Array, Dynamic Programming, State Machines

📋 Problem Statement

You are given an array prices where prices[i] is the price of a given stock on the i-th 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 1 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]
Day 0: Buy at price 1
Day 1: Sell at price 2, profit = 2-1 = 1
Day 2: Cooldown (cannot buy)
Day 3: Buy at price 0
Day 4: Sell at price 2, profit = 2-0 = 2
Total profit = 1 + 2 = 3

Example 2:

Input: prices = [1]
Output: 0
Explanation: Cannot complete any transaction with just one day.

Constraints

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

💡 Implication of constraints: - ✅ At least one day of trading (minimum array size 1) - ✅ Can buy and sell multiple times but with cooldown constraint - ✅ Must sell before buying again (hold at most one share) - ✅ After selling, must wait 1 day before buying again - ✅ Medium input size allows O(n) solution

🧒 Understanding The Problem (By Hand)

Let Me Work Example 1 Step By Step

prices = [1, 2, 3, 0, 2]
          0  1  2  3  4  (day indices)

Rules: 
  - Can buy and sell multiple times
  - Hold at most one stock at a time
  - After selling, cooldown for 1 day (cannot buy next day)

Let me trace the optimal strategy:

Day 0: Buy at price 1 (holding stock)
Day 1: Sell at price 2 (profit = 1, enter cooldown)
Day 2: Cooldown (cannot buy, must wait)
Day 3: Buy at price 0 (holding stock again)
Day 4: Sell at price 2 (profit = 2)

Total profit = 1 + 2 = 3 ✓

Alternative strategies:
Strategy A: Buy day 0 (1), Sell day 2 (3) → profit = 2, cooldown day 3, no more good buys → Total = 2
Strategy B: Buy day 3 (0), Sell day 4 (2) → profit = 2, but missed earlier opportunity → Total = 2

What Did I Learn?

Key Insight 1: Cooldown creates a 3-state system instead of 2-state
Key Insight 2: After selling, we must track the cooldown state
Key Insight 3: From cooldown, we can only go to "can buy" state

State Transitions:
  Can Buy → Buy Stock (go to Holding)
  Can Buy → Skip (stay in Can Buy)
  Holding → Sell Stock (go to Cooldown)  
  Holding → Keep Holding (stay in Holding)
  Cooldown → Must wait (go to Can Buy)

🚫 Approach 1 — Brute Force Recursion

Thought Process

At each day, we have three possible states: - Can Buy (state = 0): We don't currently hold any stock and can buy - Holding (state = 1): We currently hold a stock and can sell - Cooldown (state = 2): We just sold and must wait 1 day

For each state, we have different choices: 1. When we can buy: Skip buying OR buy the stock 2. When we're holding: Skip selling OR sell the stock (enter cooldown) 3. When in cooldown: Must skip (automatically go to can buy next day)

We'll explore all possibilities recursively and take the maximum profit.

Algorithm / Pseudocode

solve(day, state, prices):
    if day == n: return 0  // No more days

    if state == 0:  // Can buy
        skip = solve(day+1, 0, prices)    // Skip buying, stay ready
        buy = -prices[day] + solve(day+1, 1, prices)  // Buy, now holding
        return max(skip, buy)

    elif state == 1:  // Holding stock
        skip = solve(day+1, 1, prices)   // Keep holding
        sell = prices[day] + solve(day+1, 2, prices)   // Sell, enter cooldown
        return max(skip, sell)

    else:  // state == 2, Cooldown
        return solve(day+1, 0, prices)   // Must wait, then can buy

Code

class Solution {
    // Function to find the maximum profit earned using recursion
    private int solve(int day, int state, int n, int[] prices) {
        // Base case: reached end of array
        if (day == n) {
            return 0;
        }

        int profit = 0;

        // State 0: Can buy stock
        if (state == 0) { 
            profit = Math.max(
                solve(day + 1, 0, n, prices),                    // Skip buying
                -prices[day] + solve(day + 1, 1, n, prices)      // Buy stock
            );
        }

        // State 1: Holding stock
        else if (state == 1) { 
            profit = Math.max(
                solve(day + 1, 1, n, prices),                    // Keep holding
                prices[day] + solve(day + 1, 2, n, prices)       // Sell stock (enter cooldown)
            );
        }

        // State 2: Cooldown
        else { 
            profit = solve(day + 1, 0, n, prices);              // Must wait, then can buy
        }

        return profit;
    }

    public int maxProfit(int[] prices) {
        int n = prices.length;
        if (n == 0) return 0;

        // Start at day 0, can buy (state 0)
        return solve(0, 0, n, prices);
    }
}

Dry Run

For prices = [1, 2, 3]:

solve(0, 0) → Can buy at day 0 (price=1)
├─ Skip: solve(1, 0)
│  ├─ Skip: solve(2, 0) 
│  │  ├─ Skip: solve(3, 0) → 0
│  │  └─ Buy: -3 + solve(3, 1) → -3 + 0 = -3
│  │  → max(0, -3) = 0
│  └─ Buy: -2 + solve(2, 1)
│     ├─ Hold: solve(3, 1) → 0  
│     └─ Sell: 3 + solve(3, 2) → 3 + 0 = 3
│     → max(0, 3) = 3
│  → max(0, -2+3) = 1
└─ Buy: -1 + solve(1, 1) 
   ├─ Hold: solve(2, 1) → max(solve(3,1), 3+solve(3,2)) = max(0, 3) = 3
   └─ Sell: 2 + solve(2, 2) → 2 + solve(3, 0) → 2 + 0 = 2
   → max(3, 2) = 3
→ max(1, -1+3) = max(1, 2) = 2

This creates exponential tree with overlapping subproblems...

Time & Space Complexity

⏰ Time: O(3^n) - Each call branches into 2-3 possibilities
💾 Space: O(n) - Recursion depth

This approach has exponential time complexity due to overlapping subproblems.

⚠️ Approach 2 — Memoization (Top-Down DP)

Thought Process (How did you arrive here)

🎯 The Optimization Insight:

Brute Force Problem: Same subproblems solved repeatedly
                    ↓
Key Observation: solve(day, state) returns same value every time
                    ↓
Solution: Store results in dp[day][state] table
         Check table before computing
                    ↓
Trade-off: Use O(n × 3) space to achieve O(n × 3) time

Algorithm / Pseudocode

1. Create dp[n][3] initialized with -1
2. Before computing solve(day, state), check if dp[day][state] != -1
3. If already computed, return dp[day][state]
4. Otherwise, compute, store, and return: dp[day][state] = result

Code

class Solution {
    // Function to find the maximum profit earned using memoization
    private int solve(int day, int state, int n, int[] prices, int[][] dp) {
        // Base case
        if (day == n) {
            return 0;
        }

        // Check if already computed
        if (dp[day][state] != -1) {
            return dp[day][state];
        }

        int profit = 0;

        // State 0: Can buy
        if (state == 0) { 
            profit = Math.max(
                solve(day + 1, 0, n, prices, dp),
                -prices[day] + solve(day + 1, 1, n, prices, dp)
            );
        }

        // State 1: Holding
        else if (state == 1) { 
            profit = Math.max(
                solve(day + 1, 1, n, prices, dp),
                prices[day] + solve(day + 1, 2, n, prices, dp)
            );
        }

        // State 2: Cooldown
        else { 
            profit = solve(day + 1, 0, n, prices, dp);
        }

        // Store and return result
        return dp[day][state] = profit;
    }

    public int maxProfit(int[] prices) {
        int n = prices.length;
        if (n == 0) return 0;

        // Initialize DP table
        int[][] dp = new int[n][3];
        for (int[] row : dp) {
            Arrays.fill(row, -1);
        }

        return solve(0, 0, n, prices, dp);
    }
}

Dry Run

For prices = [1, 2, 3, 0, 2]:

DP Table Construction:
═══════════════════════════════════════════════════════════════════
Day   Price   dp[day][0]   dp[day][1]   dp[day][2]   Meaning
                (canBuy)     (holding)    (cooldown)
═══════════════════════════════════════════════════════════════════
 0      1         3           2           X        Best from day 0
 1      2         3           3           2        Best from day 1  
 2      3         2           2           2        Best from day 2
 3      0         2           2           0        Best from day 3
 4      2         0           0           0        Best from day 4

Answer: dp[0][0] = 3

Time & Space Complexity

⏰ Time: O(n × 3) = O(n) - Each state computed once
💾 Space: O(n × 3) + O(n) = O(n) - DP table + recursion stack

✅ Approach 3 — Tabulation (Bottom-Up DP)

Thought Process (How did you arrive here)

💡 Eliminating Recursion:

Memoization Limitation: Still uses O(n) recursion stack
                       ↓
Key Insight: Fill DP table iteratively from end to start
            Eliminate recursion overhead
                       ↓
Algorithm: Start from base case, work towards answer
           dp[day] depends only on dp[day+1]
                       ↓
Result: Same time complexity, cleaner implementation

Base Case Derivation (Step by Step)

🤔 The Big Question: What happens when we've processed all days? What should our DP table return?

Let's think about this intuitively:

Step 1: Understanding the "No More Days" Scenario

Imagine: prices = [1, 2, 3] and we're at day 3 (beyond last day)

Question: If there are no more days left, what's the maximum profit?
Answer: 0! We can't make any more moves.

Why 0 regardless of our state?
  - If we can buy: No days left to buy and sell → profit = 0
  - If we're holding: No days left to sell → stuck with stock → profit = 0
  - If we're in cooldown: No days left to utilize → profit = 0

Conclusion: dp[n][0] = dp[n][1] = dp[n][2] = 0

Step 2: Base Case in Code

From our recursive logic:
  if (day == n) return 0;  ← BASE CASE!

This means:
  solve(n, 0) = 0  (no more days to trade)
  solve(n, 1) = 0  (no more days to trade)
  solve(n, 2) = 0  (no more days to trade)

In tabulation terms:
  dp[n][0] = dp[n][1] = dp[n][2] = 0

Algorithm / Pseudocode

1. Create dp[n+1][3] (extra row for base case)
2. Initialize base case: dp[n][0] = dp[n][1] = dp[n][2] = 0
3. Iterate from day (n-1) down to 0:
   For each state (0 to 2):
     Apply same transition logic as recursion:
     - State 0: max(skip, buy)
     - State 1: max(hold, sell → cooldown)
     - State 2: must go to state 0
4. Return dp[0][0] (max profit starting from day 0, able to buy)

Code

class Solution {
    public int maxProfit(int[] prices) {
        int n = prices.length;
        if (n == 0) return 0;

        // Create DP table with base case
        int[][] dp = new int[n + 1][3];

        // Base condition: no profit at the end
        dp[n][0] = dp[n][1] = dp[n][2] = 0;

        // Fill table from bottom-up
        for (int day = n - 1; day >= 0; day--) {
            // State 0: Can buy
            dp[day][0] = Math.max(
                dp[day + 1][0],                         // Skip buying
                -prices[day] + dp[day + 1][1]           // Buy stock
            );

            // State 1: Holding
            dp[day][1] = Math.max(
                dp[day + 1][1],                         // Keep holding
                prices[day] + dp[day + 1][2]            // Sell stock (enter cooldown)
            );

            // State 2: Cooldown
            dp[day][2] = dp[day + 1][0];                // Must wait, then can buy
        }

        return dp[0][0];
    }
}

Dry Run

Building DP table for prices = [1, 2, 3, 0, 2]:

Bottom-Up Table Construction:
═══════════════════════════════════════════════════════════════════

Step 1: Base case (day 5)
dp[5][0] = dp[5][1] = dp[5][2] = 0

Step 2: Day 4 (price = 2)
State 0: max(dp[5][0], -2 + dp[5][1]) = max(0, -2+0) = 0
State 1: max(dp[5][1], 2 + dp[5][2]) = max(0, 2+0) = 2
State 2: dp[5][0] = 0

Step 3: Day 3 (price = 0)  
State 0: max(dp[4][0], -0 + dp[4][1]) = max(0, 0+2) = 2
State 1: max(dp[4][1], 0 + dp[4][2]) = max(2, 0+0) = 2
State 2: dp[4][0] = 0

Step 4: Day 2 (price = 3)
State 0: max(dp[3][0], -3 + dp[3][1]) = max(2, -3+2) = 2
State 1: max(dp[3][1], 3 + dp[3][2]) = max(2, 3+0) = 3
State 2: dp[3][0] = 2

Step 5: Day 1 (price = 2)
State 0: max(dp[2][0], -2 + dp[2][1]) = max(2, -2+3) = 2
State 1: max(dp[2][1], 2 + dp[2][2]) = max(3, 2+2) = 4
State 2: dp[2][0] = 2

Step 6: Day 0 (price = 1)
State 0: max(dp[1][0], -1 + dp[1][1]) = max(2, -1+4) = 3 ← ANSWER
State 1: max(dp[1][1], 1 + dp[1][2]) = max(4, 1+2) = 4
State 2: dp[1][0] = 2

═══════════════════════════════════════════════════════════════════
Final Answer: dp[0][0] = 3

Time & Space Complexity

⏰ Time: O(n × 3) = O(n)
💾 Space: O(n × 3) = O(n)

🎯 Approach 4 — Space Optimized (Final)

Thought Process (How did you arrive here)

🚀 Ultimate Optimization:

Tabulation Observation: dp[day] only depends on dp[day+1]
                       ↓
Space Insight: Don't need entire n×3 table
               Only need "current" and "next" day values
                       ↓
Optimization: Use 6 variables instead of 2D array
              next_canBuy, next_holding, next_cooldown
              curr_canBuy, curr_holding, curr_cooldown
                       ↓
Result: Same time O(n), space reduced to O(1)

Algorithm / Pseudocode

1. Initialize: next_canBuy=0, next_holding=0, next_cooldown=0 (base case)
2. For day from n-1 down to 0:
   - curr_canBuy = max(next_canBuy, -price + next_holding)
   - curr_holding = max(next_holding, price + next_cooldown)
   - curr_cooldown = next_canBuy
   - Update next values for next iteration
3. Return curr_canBuy

Code

class Solution {
    public int maxProfit(int[] prices) {
        int n = prices.length;
        if (n == 0) return 0;

        // Space optimization: only need next day's values
        int next_canBuy = 0, next_holding = 0, next_cooldown = 0;
        int curr_canBuy = 0, curr_holding = 0, curr_cooldown = 0;

        for (int day = n - 1; day >= 0; day--) {
            // Current day: can buy
            curr_canBuy = Math.max(next_canBuy, -prices[day] + next_holding);

            // Current day: holding
            curr_holding = Math.max(next_holding, prices[day] + next_cooldown);

            // Current day: cooldown  
            curr_cooldown = next_canBuy;

            // Update for next iteration
            next_canBuy = curr_canBuy;
            next_holding = curr_holding;
            next_cooldown = curr_cooldown;
        }

        return curr_canBuy;
    }
}

Time & Space Complexity

⏰ Time: O(n)
💾 Space: O(1)

Quick Revision

Quick Implementation

public class Solution {
  public int maxProfit(int[] a) {
    // return recursive(a, 0, 0);

    // Integer[][] memo = new Integer[a.length + 1][3];
    // return topDown(a, 0, 0, memo);

    return bottomUp(a);
  }

  private int bottomUp(int[] a) {
    int n = a.length;
    int[][] dp = new int[n + 1][3];

    for (int i = n - 1; i >= 0; i--) {
      for (int state = 0; state <= 2; state++) {

        int profit = 0;
        if (state == 0) {
          int skipIt = dp[i + 1][0];
          int buyIt = -a[i] + dp[i + 1][1];

          profit = Math.max(skipIt, buyIt);
        } else if (state == 1) {
          int skipIt = dp[i + 1][1];
          int sellIt = a[i] + dp[i + 1][2];

          profit = Math.max(skipIt, sellIt);
        } else {
          profit = dp[i + 1][0];
        }

        dp[i][state] = profit;

      }
    }

    return dp[0][0];
  }

  private int topDown(int[] a, int index, int state, Integer[][] memo) {
    if (index == a.length) {
      return 0;
    }

    if (memo[index][state] != null) {
      return memo[index][state];
    }

    int profit = 0;
    if (state == 0) {
      int skipIt = topDown(a, index + 1, 0, memo);
      int buyIt = -a[index] + topDown(a, index + 1, 1, memo);

      profit = Math.max(skipIt, buyIt);
    } else if (state == 1) {
      int skipIt = topDown(a, index + 1, 1, memo);
      // GOTCHA: state moved to COOLDOWN
      int sellIt = a[index] + topDown(a, index + 1, 2, memo);

      profit = Math.max(skipIt, sellIt);
    } else {
      profit = topDown(a, index + 1, 0, memo);
    }

    return memo[index][state] = profit;
  }

  private int recursive(int[] a, int index, int state) {
    // step 2: base case
    if (index == a.length) {
      return 0;
    }

    int profit = 0;
    // step 1: state = 0 or 1 or 2
    // 0 means we can buy (not mandatory. can skip also)
    // 1 means we can sell (not mandatory. can skip also)
    // 2 means we can do nothing (no buy or sell on the current day)
    if (state == 0) {
      int skipIt = recursive(a, index + 1, 0);
      int buyIt = -a[index] + recursive(a, index + 1, 1);

      profit = Math.max(skipIt, buyIt);
    } else if (state == 1) {
      int skipIt = recursive(a, index + 1, 1);
      // GOTCHA: state moved to COOLDOWN
      int sellIt = a[index] + recursive(a, index + 1, 2);

      profit = Math.max(skipIt, sellIt);
    } else {
      profit = recursive(a, index + 1, 0);
    }

    return profit;
  }

  public static void main(String[] args) {
    Solution s = new Solution();
    System.out.println(s.maxProfit(new int[] { 1, 2, 3, 0, 2 }) == 3);
  }
}

📝 Edge Cases & Testing

Edge Cases

Single day: [5] → Return 0 (can't buy and sell for profit)
Two days increasing: [1, 5] → Return 4 (buy day 1, sell day 2, no cooldown issue)
Two days decreasing: [5, 1] → Return 0 (don't buy at all)
Three days with cooldown: [1, 2, 1] → Return 1 (buy day 1, sell day 2, cooldown day 3)
All same prices: [3, 3, 3] → Return 0 (no profit opportunity)
Monotonic increasing: [1, 2, 3, 4] → Return 3 (buy first, sell last, cooldown doesn't matter)
Monotonic decreasing: [4, 3, 2, 1] → Return 0 (never buy)

Test Suite

Basic Tests:

assertEquals(3, maxProfit([1,2,3,0,2]));      // Example 1
assertEquals(0, maxProfit([1]));              // Example 2
assertEquals(2, maxProfit([2,1,4]));          // Simple profitable case

Edge Tests:

assertEquals(0, maxProfit([5]));              // Single element
assertEquals(4, maxProfit([1,5]));            // Two elements up
assertEquals(0, maxProfit([5,1]));            // Two elements down
assertEquals(0, maxProfit([3,3,3,3]));        // All same
assertEquals(1, maxProfit([1,2,1]));          // Cooldown effect test

Cooldown Specific Tests:

assertEquals(3, maxProfit([1,2,1,3]));        // Buy-sell-cooldown-buy-sell
assertEquals(2, maxProfit([2,1,2,1,2]));      // Multiple small transactions
assertEquals(4, maxProfit([1,2,4]));          // Hold through vs sell early

Stress Tests:

// Large input with max profit
int[] largeUp = new int[5000];
for(int i = 0; i < 5000; i++) largeUp[i] = i;
// Expected: can't get full n-1 profit due to cooldowns
assertTrue(maxProfit(largeUp) > 0 && maxProfit(largeUp) < 4999);

// Large input with no profit  
int[] largeDown = new int[5000];
for(int i = 0; i < 5000; i++) largeDown[i] = 5000-i;
assertEquals(0, maxProfit(largeDown));

Key Insights

State modeling: Three states perfectly capture the cooldown constraint
State transitions: Cooldown creates forced transition (cooldown → can buy)
DP progression: Recursion → memoization → tabulation → space optimization
Cooldown impact: Prevents immediate re-buying, requiring strategic timing
Pattern foundation: This 3-state pattern handles time-based constraints elegantly

Comparison with Other Stock Problems

Aspect	Stock II (Unlimited)	Stock with Cooldown
State Space	2D: [day][canBuy]	2D: [day][state] (3 states)
Complexity	O(n) time, O(1) space	O(n) time, O(1) space
Key Constraint	Hold at most 1 stock	Hold at most 1 stock + cooldown
Greedy Solution	✅ Sum positive diffs	❌ Need DP
States	Can Buy, Holding	Can Buy, Holding, Cooldown

Advanced Patterns

Forced transitions: Cooldown state has no choice (must wait)
Time dependencies: Current action affects future availability
Strategic timing: May hold longer to avoid cooldown penalties
State explosion: Each constraint adds dimensions to state space

The space-optimized solution achieves optimal O(n) time and O(1) space complexity while elegantly handling the cooldown constraint through a 3-state machine!