leetcode 123. 买卖股票的最佳时机 III Python

题目

给定一个数组，它的第 i 个元素是一支给定的股票在第 i 天的价格。

设计一个算法来计算你所能获取的最大利润。你最多可以完成 两笔 交易。

注意: 你不能同时参与多笔交易（你必须在再次购买前出售掉之前的股票）,地址。
示例：

输入: [3,3,5,0,0,3,1,4]
输出: 6
解释: 在第 4 天（股票价格 = 0）的时候买入，在第 6 天（股票价格 = 3）的时候卖出，这笔交易所能获得利润 = 3-0 = 3 。
     随后，在第 7 天（股票价格 = 1）的时候买入，在第 8 天 （股票价格 = 4）的时候卖出，这笔交易所能获得利润 = 4-1 = 3 。

解题思路

最佳买卖股票时机含冷冻期
类似于含冷冻期的买卖股票，也是一个状态机。本次买卖股票共有四种状态：第一次买卖股票，第一次卖出股票后，第二次买入股票，第二次卖出股票后分别设位S1,S2,S3,S4，并且对应的状态转移方程：

s1 = max(s1, -prices[i])
s2 = max(s2, s1 + prices[i])
s3 = max(s3, s2 - prices[i])
s4 = max(s4, s3 + prices[i])                

Python实现

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        size = len(prices)
        if size == 0:
            return 0
        
        # s1是假设第一天就买入股票，s2，s3，s4则分别设为最小值
        s1 = -prices[0]
        s2, s3, s4 = float('-inf'), float('-inf'), float('-inf')
        for i in range(size):
            s1 = max(s1, -prices[i])
            s2 = max(s2, s1 + prices[i])
            s3 = max(s3, s2 - prices[i])
            s4 = max(s4, s3 + prices[i])                
                
        return max(0, s4)