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

The link to the problem: https://leetcode.com/problems/best-time-to-buy-and-sell-stock-with-cooldown/

Say you have an array for which the ith element is the price of a given stock on Day i. Design an algorithm to find the maximum profit. You may complete as many transactions as you like (ie, buy one and sell one share of the stock multiple times) with the following restrictions:

• You may not engage in multiple transactions at the same time (ie, you must sell the stock before you buy again).
• After you sell your stock, you cannot buy stock on the next day. (ie, cooldown 1 day)

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

Thoughts:

So far we have practiced several stock problems like Leetcode — 122. Best Time to Buy and Sell Stock II or Leetcode — 123. Best Time to Buy and Sell Stock III. We have a new condition of buying and selling stock, this time, we need to take care of the cooldown time.

We can still use dynamic programming to solve this problem since we need to decide today’s action depends on the actions in previous days.

The easiest and the most straightforward way to solve this problem is using brute force. We use an array to store the max profit we can earn till Day i.

On each day, we try to see which action we choose will help us gain the max profit: cooldown OR sell.
⚠️The reason why we don’t consider buying today is that the profit will be lower if we choose to buy. We only decide the buying day once we decide the selling day.

1. If we choose to cool down today, that means no extra profit we can earn today. Hence, the max profit we can earn today is the same as yesterday.
2. If we choose to sell the stock today, then it’s time for us to decide the buying day. We need to traverse all the previous days to see when is our best day to buy the stock. Also, if we decide to buy on Day j and sell it today, we can also take the max profit we can earn two days before Day j.
   👷The reason why it’s 2 days before Day j: if we choose to buy on Day j, that means Day (j-1) cannot be selling day nor buying day, it must be cooldown day. Hence, it is safe to take the max profit we earn on two days before Day j since as long as we are sure Day (j-1) is cooldown day, we don’t need to worry or care about whether we sell or cooldown on Day (j-2).

Below is the 29ms code that beats 7% —

public int maxProfit(int[] prices) {
        int n = prices.length;
        if(n<=1)
            return 0;
        int[] dp = new int[n];
        if(prices[1]>prices[0])
            dp[1] = prices[1]-prices[0];
        for(int i=2;i<n;i++)
        {
            int maxProfit = dp[i-1]; //cooldown
            //find buy day
            for(int j=0;j<i;j++)
            {
                if(prices[j]>prices[i])
                    continue;
                int preProfit = j>1? dp[j-2]:0; 
                maxProfit = Math.max(maxProfit,
                                    preProfit+prices[i]-prices[j]);
            }
            dp[i] = maxProfit;
        }
        return dp[n-1];
    }

Time Complexity: O(n²) ; Space Complexity: O(n).

Can we solve the problem in O(n)?

Code Logic:

Since it requires an extra O(n) time to find the buying day after we decided the selling day, is it possible that we can skip this step?

This time, we will consider three different actions for each day. We will consider the different profits if we buy, sell, or cooldown on Day i. We will use 3 arrays to track all the possible max profits if our last action is to buy, sell, or cooldown TILL Day i.

We always need to be careful about what action we took TILL Day (i-1) if we choose to do certain actions ON Day i. We always need to check if we decide to buy/sell/cooldown today, can we earn more profit than yesterday, if not, we will give up and remain the max profit we met so far.

🚧In the buy/sell/cooldown array, we will separately record its max profit we can get so far with the condition that the last action is to buy/sell/cooldown.

Three cases needed to be considered on Day i:

1. If we decide to buy today, that means the last action till yesterday MUST BE cooldown since we cannot buy right after sell. Then, we need to compare two profits: “if we cooldown yesterday and buy today” vs. “the max profit yesterday when the last action is to buy”.
2. If we decide to sell today, that means the last action till yesterday MUST BE buy since we need to buy if we want to sell. Then we compare two profits: “if we sell the stock based on our last buy” vs. “the max profit yesterday when the last action is to sell”.
3. If we decide to cool down today, that means the last action till yesterday MUST BE sell OR cooldown. We compare two profits: “the max profit yesterday when the last action is to sell” vs. “the max profit yesterday when the last action is cooldown”.

Below is the 1ms code that beats 55% —

public int maxProfit(int[] prices) {
        int n = prices.length;
        if(n<=1)
            return 0;
        int[] buy = new int[n];
        int[] sell = new int[n];
        int[] cooldown = new int[n];
        
        //no choice on day 0
        buy[0] = -prices[0];
        for(int i=1;i<n;i++)
        {
            buy[i] = Math.max(buy[i-1], cooldown[i-1]-prices[i]);
            sell[i] = Math.max(sell[i-1], prices[i]+buy[i-1]);
            cooldown[i] = Math.max(cooldown[i-1], sell[i-1]);
        }
        return Math.max(sell[n-1],cooldown[n-1]);

Time Complexity: O(n); Space Complexity: O(n)

Based on exactly the same code logic, we can also get rid of 3 arrays and convert them to 3 variables since we only need the (i-1)th value in each array.

Below is the 0ms code that beats 100% —

public int maxProfit(int[] prices) {
        int n = prices.length;
        if(n<=1)
            return 0;
        int buy = -prices[0];
        int sell = 0;
        int cooldown = 0; 
        for(int i=1;i<n;i++)
        {
            int pre_buy = buy, pre_sell = sell;
            int pre_cooldown = cooldown;
            buy = Math.max(pre_buy, pre_cooldown-prices[i]);
            sell = Math.max(pre_sell, prices[i]+pre_buy);
            cooldown = Math.max(pre_cooldown,pre_sell);
        }
        return Math.max(sell,cooldown);
    }

Time Complexity: O(n); Space Complexity: O(1)