Leetcode 309. Best Time to Buy and Sell Stock with Cooldown (Python)

Related Topic

Dynamic-Programming.

Description

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)

Sample I/O

Example 1

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

Methodology

This question solved by Dynamic Programming.

This question is hard. I found a great solution from some other sits.

Based on the question, each element should have three status which are buy, sell and cool down. I was struggled here for quite long time to try to make dp list. However, you can actually sort of convert these three status to 2 status. Hold stock or unhold stock:

• Hold stock: At the end of today, the maxmium profit if you hold stock and these can be two conditions:
  1. You bought the stock today
  2. You did not sell the stock that you bought before
• Unhold stock: At the end of today, the maxmium profit if you unhold stock and these can be two conditions as well:
  1. You sold the stock today
  2. You did not buy any stock after you sold last stock

The answer should be when you not hold stock, the maxmium profit.

Find the base case:

Create two dp list one is called hold, the other called unhold. each element represent the maxmium profit at the end of ith day. Init the first element of hold to -price[0] Init the first element of unhold to 0

Find the pattern:

There should be two patterns one for hold stock and the other one for unhold stock:

• hold stock: max(you bought the stock today, you bought the stock before but not sold it yet)
• unhold stock: max(you sold the stock today, you did not buy any stock after sold last stock)

Answer:

last element of unhold stock.

Code

def maxProfit(self, prices: List[int]) -> int:
        hold = [0]*L #hold[i]: the ith day profit when you hold stock
        unhold = [0]*L #unhold[i]: the ith day profit when you not hold stock
        
        hold[0] = -prices[0]
        unhold[0] = 0
        
        for i in range(1,L):
            hold[i] = max(unhold[i-2]-prices[i], hold[i-1])#max(today buy in, last buy in)
            unhold[i] = max(hold[i-1]+prices[i], unhold[i-1])#max(today sell out, not sell out)
        return unhold[L-1]

BigO

Then iterate both dp lists, it will cost O(2L). generally it is O(L)