309. Best Time to Buy and Sell Stock with Cooldown - Leetcode Solution

Code Implementation

class Solution:
    def maxProfit(self, prices):
        if not prices:
            return 0
        n = len(prices)
        hold = [0] * n
        sold = [0] * n
        rest = [0] * n

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

        for i in range(1, n):
            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])

        return max(sold[-1], rest[-1])

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        if (prices.empty()) return 0;
        int n = prices.size();
        vector<int> hold(n, 0), sold(n, 0), rest(n, 0);
        hold[0] = -prices[0];
        sold[0] = 0;
        rest[0] = 0;
        for (int i = 1; i < n; ++i) {
            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]);
        }
        return max(sold[n-1], rest[n-1]);
    }
};

class Solution {
    public int maxProfit(int[] prices) {
        if (prices.length == 0) return 0;
        int n = prices.length;
        int[] hold = new int[n];
        int[] sold = new int[n];
        int[] rest = new int[n];
        hold[0] = -prices[0];
        sold[0] = 0;
        rest[0] = 0;
        for (int i = 1; i < n; i++) {
            hold[i] = Math.max(hold[i-1], rest[i-1] - prices[i]);
            sold[i] = hold[i-1] + prices[i];
            rest[i] = Math.max(rest[i-1], sold[i-1]);
        }
        return Math.max(sold[n-1], rest[n-1]);
    }
}

var maxProfit = function(prices) {
    if (prices.length === 0) return 0;
    const n = prices.length;
    let hold = new Array(n).fill(0);
    let sold = new Array(n).fill(0);
    let rest = new Array(n).fill(0);
    hold[0] = -prices[0];
    sold[0] = 0;
    rest[0] = 0;
    for (let i = 1; i < n; i++) {
        hold[i] = Math.max(hold[i-1], rest[i-1] - prices[i]);
        sold[i] = hold[i-1] + prices[i];
        rest[i] = Math.max(rest[i-1], sold[i-1]);
    }
    return Math.max(sold[n-1], rest[n-1]);
};

Problem Description

You are given an array prices, where prices[i] is the price of a given stock on day i. You want to maximize your profit by choosing a sequence of buy and sell operations, but with these constraints:

You may complete as many transactions as you like (buy one and sell one share of the stock multiple times).
After you sell a stock, you cannot buy stock on the next day (i.e., there is a cooldown of one day).
You may not hold more than one share at a time (i.e., you must sell before you buy again).

Your task is to return the maximum profit you can achieve given these rules. The input is the prices array, and you must not reuse elements or violate the cooldown constraint.

Thought Process

At first glance, this problem seems similar to the classic "Best Time to Buy and Sell Stock" problems. However, the cooldown adds a twist: after selling, you must wait one day before buying again. Brute-force approaches might try every possible sequence of buys and sells, but this quickly becomes infeasible due to exponential growth in possibilities.

To optimize, we notice that at each day, our decision depends on our current "state": are we holding a stock, resting, or have we just sold? This suggests a state-machine or dynamic programming approach, where we calculate the best profit possible for each state on each day, based on transitions from previous days.

Solution Approach

The optimal approach uses dynamic programming with three states for each day:

hold[i]: The maximum profit if we are holding a stock on day i.
sold[i]: The maximum profit if we have just sold a stock on day i.
rest[i]: The maximum profit if we are in cooldown or resting on day i (i.e., not holding and not just sold).

The transitions are:

hold[i] = max(hold[i-1], rest[i-1] - prices[i])
(Either we were already holding, or we buy today after a rest.)
sold[i] = hold[i-1] + prices[i]
(We sell today, so yesterday we had to be holding.)
rest[i] = max(rest[i-1], sold[i-1])
(Either we continue to rest, or we enter rest after selling yesterday.)

The answer is the maximum of sold[n-1] and rest[n-1] on the last day, since we can't end with a "hold" (that would mean we still own stock).

We initialize:

hold[0] = -prices[0] (buy on day 0)
sold[0] = 0
rest[0] = 0

We iterate through days 1 to n-1, updating each state according to the above transitions.

Example Walkthrough

Let's use the input prices = [1,2,3,0,2].

Day 0: Buy at 1.
hold = -1; sold = 0; rest = 0
Day 1: Price = 2
hold = max(-1, 0-2) = -1; (keep holding)
sold = -1+2 = 1; (sell today)
rest = max(0,0) = 0
Day 2: Price = 3
hold = max(-1,0-3) = -1
sold = -1+3 = 2
rest = max(0,1) = 1
Day 3: Price = 0
hold = max(-1,2-0) = 2
sold = -1+0 = -1
rest = max(1,2) = 2
Day 4: Price = 2
hold = max(2,2-2) = 2
sold = 2+2 = 4
rest = max(2,-1) = 2

The maximum profit is max(sold[4], rest[4]) = max(4,2) = 4.
The optimal sequence: buy at 1, sell at 3, cooldown, buy at 0, sell at 2.

Time and Space Complexity

Brute-force approach: Exponential time, as each day you have two choices (buy/sell or not), leading to O(2^n) possibilities.
Optimized DP approach: O(n) time and O(n) space, since we process each day once and maintain three arrays of size n. This can be further improved to O(1) space by only keeping track of previous day's states.

The O(n) time comes from the single pass through the array, and the O(n) space from the state arrays.

Summary

This problem is elegantly solved with dynamic programming by modeling the three possible states (hold, sold, rest) at each day and updating them based on previous values. The cooldown constraint is naturally handled by the "rest" state. By thinking in terms of state transitions rather than brute-forcing all transactions, we achieve an efficient and clear solution.

AlgoMap.io

Want Help Cracking FAANG?