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

Question

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)

现在给你一个数组，第i个元素代表第i天的股票价格。设计出一个算法去找到最大的收益。允许多次交易，但是你必须在第二次买股票前卖掉股票，而且当你卖掉股票，第二天你不被允许买股票（冷却时间）

Example

prices = [1, 2, 3, 0, 2]
maxProfit = 3
transactions = [buy, sell, cooldown, buy, sell]

Solution

• 动态规划解。思考问题仍然从如何找到第i天和前一天的最大收益的关系出发。我们可以发现每一天都有三种状态，即买（buy），卖（sell），不买不卖（rest），那么其实我们可以分别记录第i天buy, sell, rest的最大收益记录下来，然后根据它们之间的关系更新。这个思路完全没有问题，但是在实际操作的时候我们会发现，rest其实也要分两种情况。

  • 第一种是sell后，buy前的rest，因为sell后不能马上buy，所以如果第i天是rest_before_buy，那么第i - 1天只能是rest_before_buy本身或者sell
  • 第二种是buy后，sell前的rest，如果第i天是rest_after_buy，那么第i - 1天只能是rest_after_buy或者buy

  所以这两种rest并不是同一种，因此其实每一天有四种状态，即rest_before_buy, buy, rest_after_buysell, sell，关系如下图

  

  • 关系式为：

    buy[i] = rest_before_buy[i - 1] - prices[i]

    rest_after_buy[i] = max(buy[i - 1], rest_after_buy[i - 1])

    sell[i] = max(buy[i - 1] + prices[i], rest_after_buy[i - 1] + prices[i])

    rest_before_buy[i] = max(sell[i - 1], rest_before_buy[i - 1])

  • 初始化为：

    buy[0] = -prices[0] # 第一天买进花费prices[0]

    rest_after_buy[0] = -prices[0] # 因为是不买不卖，所以和buy[0]相同

    sell[0] = - (prices[0] + 1) # 第一天不可能卖的，所以设一个非常小的值

    rest_before_buy = 0 # 第一天如果不买不卖，花费0

  • 代码

    class Solution(object):
        def maxProfit(self, prices):
            """
            :type prices: List[int]
            :rtype: int
            """
            if len(prices) == 0: return 0
            sell = [0] * len(prices)
            buy = [0] * len(prices)
            rest_before_buy = [0] * len(prices)
            rest_after_buy = [0] * len(prices)
    
            # 初始化
            buy[0] = - prices[0]
            sell[0] = - (prices[0] + 1)
            rest_after_buy[0] = - prices[0]
    
            for index_p in range(1, len(prices)):
                buy[index_p] = rest_before_buy[index_p - 1] - prices[index_p]
                rest_after_buy[index_p] = max(buy[index_p - 1], rest_after_buy[index_p - 1])
                sell[index_p] = max(rest_after_buy[index_p - 1] + prices[index_p], buy[index_p - 1] + prices[index_p])
                rest_before_buy[index_p] = max(sell[index_p - 1], rest_before_buy[index_p - 1])
    
            # 因为最后一天buy或者rest_after_buy（有股票没卖掉）都不可能利益最大的，所以不需比较
            return max(sell[-1], rest_before_buy[-1])