Daily LeetCode 309. Best Time to Buy and Sell Stock with Cooldown

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

Medium

问题描述

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]

思路及代码

这又是一道与股票相关的问题，与之前的买入卖出不同，这一题引入了cooldown概念，每当执行一次卖出操作后，第二天必须cooldown一天，不能够进行买入操作。在这样一个操作规则下，我们需要求出最大的利润。

我们可以用一个自动机来表示这个问题的三个状态：

三个状态的更新过程如下：

#保持s0状态，或者从cooldown状态出来
s0[i] = max(s0[i - 1], s2[i - 1]);
#保持s1状态，或者执行购入操作
s1[i] = max(s1[i - 1], s0[i - 1] - prices[i]);
#只有一个操作，s1->s2
s2[i] = s1[i - 1] + prices[i];

代码如下：

def maxProfit(prices):
    s0, s1, s2 = [0] * len(prices), [0] * len(prices), [0] * len(prices)
    s0[0], s1[0], s2[0] = 0, -prices[0], -sys.maxsize
    for i in range(1, len(prices)):
        s0[i] = max(s0[i - 1], s2[i - 1])
        s1[i] = max(s1[i - 1], s0[i - 1] - prices[i])
        s2[i] = s1[i - 1] + prices[i]

    return max(s0[-1], s2[-1])

由于每一时刻状态只和前一个时刻有关，我们可以将空间复杂度减少到O(1):

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        if len(prices) == 0: return 0
        buy, wait_for_sell, cooldown = 0, -prices[0], -sys.maxsize
        for i in range(1, len(prices)):
            pre_state_sell = wait_for_sell
            wait_for_sell = max(wait_for_sell, buy - prices[i])
            buy = max(buy, cooldown)
            cooldown = pre_state_sell + prices[i]

        return max(cooldown, buy)