Q309: 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 next day. (ie, cooldown 1 day)
Example:

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

At first, I thought about dynamic programming with the thought that:

For the first j days, the best solution is among: 1. the first j-1 day and do nothing at j; 2. bug and sell at kth and jth day, and doing something at the first k-2 days (we need k-2 since we cannot do anything at k-1th day in order to buy at kth day).

With this idea, it yields an DP algorithm with time complexity $O(N^2)$ since at each time I need to traverse to get the choice of k.

OK, how to improve? It can be done in linear time ..

How about maintaining some finer information? At jth day, all I need to decide is whether I want to buy, sell, or cool down (I have to cool down). The rule is (from previous day to current day):

1. sell -> cool down;
2. buy -> sell;
3. cool down -> buy

Of course, we can do nothing as well. So, essentially, we have 2 choices: 1. to sell (if have share); 2. to have share (by buying or not selling). But we need to know if we could buy so we also need to know the status about not having a share. So, we maintain three events: 1. sell at j day; 2. have a share at j day; 3. not have a share at j day (not by selling). And taking DP in, we want to maintain the best solution where the last state is one of the three respectively.

So, the pseudocode is:

best_sell = [ 0 ] * len(prices)
best_hold = [ 0 ] * len(prices)
best_noshare = [ 0 ] * len(prices)

# initialize
best_sell[0] = -float('inf')
best_hold[0] = -prices[0]
best_noshare[0] = 0

# iteration
for j in 1 .. len(prices) - 1 {
  best_sell[j] = best_hold[j-1] + prices[j]
  best_hold[j] = max(best_hold[j-1], best_noshare[j-1] - prices[j])
  best_noshare[j] = max(best_sell[j-1], best_noshare[j-1])
}

return max(best_sell[-1], best_noshare[-1])

The Python submission.

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
    
        if len(prices) == 0:
            return 0
        
        best_sell = [ 0 ] * len(prices)
        best_hold = [ 0 ] * len(prices)
        best_noshare = [ 0 ] * len(prices)
        
        # initialize
        best_sell[0] = -float('inf')
        best_hold[0] = -prices[0]
        best_noshare[0] = 0
        
        for j in range(1, len(prices)):
            best_sell[j] = best_hold[j-1] + prices[j]
            best_hold[j] = max(best_hold[j-1], best_noshare[j-1] - prices[j])
            best_noshare[j] = max(best_sell[j-1], best_noshare[j-1])
      
        return max(best_noshare[-1], best_sell[-1])