Best Time to Buy and Sell Stock with Cooldown

Problem: Like the previous Buy and Sell stock problems, 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 with the following two restrictions:

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

Example:

Input: [2,3,4,1,5]
Output: 5 
Explanation: transactions = [buy, sell, cooldown, buy, sell]

Approach: Like in previous problems, we are going to use DP approach here too. Let's have 2 1D Tables, say Buy where Buy[i] will tell the minimum money we can invest on buying stock till ith day and another table is Sell where Sell[i] will tell the maximum money we made till ith day by selling the stocks. Obviously Sell[n] will be our answer.

Here is how we can fill these tables:

1. Buy[0] = - Prices[0] and Buy[i] = MAX(Buy[i - 1], Sell[i - 2] - prices[i]). How? Here is the explanation:
• At a particular day i, either we don't do anything i.e. Buy[i - 1]. Take the previous day amount as we did not buy.
• Or we buy but to buy we need to serve the one day cooldown. That's why we will take Sell[i - 2]. Now if we buy we are spending Sell[i - 2] - prices[i] money. 
• We take the maximum of above these two numbers (hint: negative numbers) and assign it to Buy[i].
2. Sell[i] = Max(Sell[i - 1], Buy[i - 1] + prices[i]). Here is the explanation why:
• We can decide not to do any thing. In that case we will just take the previous day's amount and assign that is Sell[i - 1].
• We can decide to sell. In that case we will get Buy[i - 1] + prices[i]. Buy[i - 1] will give the minimum amount spent on buying till (i - 1)th day. 
• We take the maximum of above these two numbers and assign it to Sell[i].

We need to take care of some corner conditions and we are good to go. This will solve the problem in linear time which is very good but can we do something to optimize it further? It is obvious that we can't optimize time complexity as we have to touch each and every element of Prices array to decide the maximum profit. Then what we can do?

We can reduce space complexity. If you see we are using O(n) space which is not required. If we see closely at the above algorithm, we will find out that to calculate Buy[i] and Sell[i], we just need Buy[i - 1], Sell [i - 1] and Sell [i - 2]. That means we can use limited number of variables instead of maintaining two whole arrays to store these values. 

Let's say sell_1 is Sell[i - 1], sell_2 is Sell[i - 2] and buy_1 is Buy[i - 1]. For Buy[i] and Sell[i], we will use variables say currBuy and currSell respectively. Now our changed algorithm will look like: 

1. buy_1 = -Prices[0], sell_2 = 0, sell_1 = 0, currBuy = 0, currSell = 0
2. FOR i := 1 To n
• currBuy = MAX(buy_1, sell_2 - Prices[i])
• currSell = MAX(sell_1, buy_1 + Prices[i])
• sell_2 = sell_1
• sell_1 = currSell
• buy_1 = currBuy
3. Return currSell

You can see there is not much of change as such in algorithm except using variables instead of arrays. Hopefully its clear now and we can go to the implementation.

 

Implementation in C#:

        public int MaxProfitWithCoolDown(int[] prices)

        {

            if (prices?.Length <= 1)

            {

                return 0;

            }

            int buy_1 = -prices[0], sell_2 = 0, sell_1 = 0, currBuy, currSell = 0;

            for (int i = 1; i < prices.Length; ++i)

            {

                currBuy = Math.Max(buy_1, sell_2 - prices[i]);

                currSell = Math.Max(sell_1, buy_1 + prices[i]);

                sell_2 = sell_1;

                sell_1 = currSell;

                buy_1 = currBuy;

            }

            return currSell;

        }

Complexity: O(n)