题目0309:最佳买卖股票时机含冷冻期

题目描述

给定一个整数数组,其中第i个元素代表了第i天的股票价格 。​

设计一个算法计算出最大利润。在满足以下约束条件下,你可以尽可能地完成更多的交易(多次买卖一支股票):

你不能同时参与多笔交易(你必须在再次购买前出售掉之前的股票)。

卖出股票后,你无法在第二天买入股票 (即冷冻期为1天)。

示例:

输入: [1,2,3,0,2]
输出: 3 
解释: 对应的交易状态为: [买入, 卖出, 冷冻期, 买入, 卖出]

解题技巧

一种常用的方法是将「买入」和「卖出」分开进行考虑:「买入」为负收益,而「卖出」为正收益。在初入股市时,你只有「买入」的权利,只能获得负收益。而当你「买入」之后,你就有了「卖出」的权利,可以获得正收益。显然,我们需要尽可能地降低负收益而提高正收益,因此我们的目标总是将收益值最大化。因此,我们可以使用动态规划的方法,维护在股市中每一天结束后可以获得的「累计最大收益」,并以此进行状态转移,得到最终的答案。

思路与算法

我们用f[i]表示第i天结束之后的累计最大收益。根据题目描述,由于我们最多只能同时买入(持有)一支股票,并且卖出股票后有冷冻期的限制,因此我们会有三种不同的状态:

我们目前持有一支股票,对应的累计最大收益记为f[i][0];

我们目前不持有任何股票,并且处于冷冻期中,对应的累计最大收益记为f[i][1];

我们目前不持有任何股票,并且不处于冷冻期中,对应的累计最大收益记为f[i][2]。

这里的处于冷冻期指的是在第i天结束之后的状态。也就是说:如果第i天结束之后处于冷冻期,那么第i+1天无法买入股票。

如何进行状态转移呢?在第i天时,我们可以在不违反规则的前提下进行买入或者卖出操作,此时第i天的状态会从第i−1天的状态转移而来;我们也可以不进行任何操作,此时第i天的状态就等同于第i-1天的状态。那么我们分别对这三种状态进行分析:

对于f[i][0],我们目前持有的这一支股票可以是在第i−1天就已经持有的,对应的状态为f[i−1][0];或者是第i天买入的,那么第i−1天就不能持有股票并且不处于冷冻期中,对应的状态为f[i−1][2]加上买入股票的负收益prices[i]。因此状态转移方程为:

f[i][0] = \max(f[i-1][0], f[i-1][2] - {\it prices}[i])

对于f[i][1],我们在第i天结束之后处于冷冻期的原因是在当天卖出了股票,那么说明在第i−1天时我们必须持有一支股票,对应的状态为f[i-1][0]加上卖出股票的正收益prices[i]。因此状态转移方程为:

f[i][1] = f[i-1][0] + {\it prices}[i]

对于f[i][2],我们在第i天结束之后不持有任何股票并且不处于冷冻期,说明当天没有进行任何操作,即第i-1天时不持有任何股票:如果处于冷冻期,对应的状态为f[i−1][1];如果不处于冷冻期,对应的状态为f[i−1][2]。因此状态转移方程为:

f[i][2] = \max(f[i-1][1], f[i-1][2])

这样我们就得到了所有的状态转移方程。如果一共有n天,那么最终的答案即为:\max(f[n-1][0], f[n-1][1], f[n-1][2])

注意到如果在最后一天(第n−1天)结束之后,手上仍然持有股票,那么显然是没有任何意义的。因此更加精确地,最终的答案实际上是f[n−1][1]和f[n−1][2]中的较大值,即:\max(f[n-1][1], f[n-1][2])

我们可以将第0天的情况作为动态规划中的边界条件:

\begin{cases} f[0][0] &= -{\it prices}[0] \\ f[0][1] &= 0 \\ f[0][2] &= 0 \end{cases}

在第0天时,如果持有股票,那么只能是在第0天买入的,对应负收益−prices[0];如果不持有股票,那么收益为零。注意到第0天实际上是不存在处于冷冻期的情况的,但我们仍然可以将对应的状态f[0][1]置为零,这其中的原因留给读者进行思考。

这样我们就可以从第1天开始,根据上面的状态转移方程进行进行动态规划,直到计算出第n-1天的结果。

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        if not prices:
            return 0

        n = len(prices)
        # f[i][0]: 手上持有股票的最大收益
        # f[i][1]: 手上不持有股票,并且处于冷冻期中的累计最大收益
        # f[i][2]: 手上不持有股票,并且不在冷冻期中的累计最大收益
        f = [[-prices[0], 0, 0]] + [[0] * 3 for _ in range(n - 1)]
        for i in range(1, n):
            f[i][0] = max(f[i - 1][0], f[i - 1][2] - prices[i])
            f[i][1] = f[i - 1][0] + prices[i]
            f[i][2] = max(f[i - 1][1], f[i - 1][2])

        return max(f[n - 1][1], f[n - 1][2])

空间优化:注意到上面的状态转移方程中,f[i][..]只与f[i−1][..]有关,而与f[i−2][..]及之前的所有状态都无关,因此我们不必存储这些无关的状态。也就是说,我们只需要将f[i−1][0],f[i−1][1],f[i−1][2]存放在三个变量中,通过它们计算出f[i][0],f[i][1],f[i][2]并存回对应的变量,以便于第i+1天的状态转移即可。

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        if not prices:
            return 0

        n = len(prices)
        f0, f1, f2 = -prices[0], 0, 0
        for i in range(1, n):
            newf0 = max(f0, f2 - prices[i])
            newf1 = f0 + prices[i]
            newf2 = max(f1, f2)
            f0, f1, f2 = newf0, newf1, newf2

        return max(f1, f2)

复杂度分析

时间复杂度:O(n),其中n为数组prices的长度。

空间复杂度:O(n)。我们需要3n的空间存储动态规划中的所有状态,对应的空间复杂度为O(n)。如果使用空间优化,空间复杂度可以优化至O(1)。