股票问题通解

股票系列一共有六个问题:

• 121. 买卖股票的最佳时机
• 122. 买卖股票的最佳时机 II
• 123. 买卖股票的最佳时机 III
• 188. 买卖股票的最佳时机 IV
• 309. 最佳买卖股票时机含冷冻期
• 714. 买卖股票的最佳时机含手续费

股票问题可以使用动态规划来求解, 这六个问题的状态转移方程可以抽象为同一个. 假设有$f(i,k,s)$:

• i: 表示第i天, $i\in [0,n-1]$
• k: 表示允许的最大交易次数.
• s: 表示第i天结束时, 手中持有的股票, $k\in [0,1]$

$f(i,k,s)$表示第i天结束时, 最大交易次数为k时, 手中持有股票数为s的最大利润.

那么最终问题的解就是: $f(n-1, k, 0)$, 即最后一天时, 最大交易次数为k的情况下, 手中没有股票时最大的利润.

股票问题I~IV的核心区别就在于k的限制.

同时根据, Sell和Buy是否会消耗最大交易次数k, 可以得出两种不同的状态转移方程. 两种状态转移方程的最终解是一样的.

对于每一天, 面临着三个选择:

• 休息
• 买入
• 卖出

1. 假设 买入 操作消耗最大交易次数

因为买入和卖出是成对的. 一次买入必定伴随着卖出, 因为最终结果手中不能有股票.

对于$f(i,k,0)$:

• 第i天选择休息: $f(i,k,0)=f(i-1,k,0)$
• 第i天选择卖出: $f(i,k,0)=f(i-1,k,1)+Prices[i]$, 卖出将获得利润.

对于$f(i,k,1)$:

• 选择休息: $f(i,k,1)=f(i-1,k,1)$
• 选择买入: $f(i,k,1)=f(i-1,k-1,0)-Prices[i]$, 此处买入操作要消耗一次交易次数, 同时利润减少.

由上可以得出状态转移方程:

$\begin{cases} f(i,k,0)=max(f(i-1,k,0), f(i-1,k,1)+Prices[i]) \\ f(i,k,1)=max(f(i-1,k,1), f(i-1,k-1,0)-Prices[i]) \end{cases}$

基准值(初始值):

• $f(i,0,0)$: 最大交易次数为0, 持有0份股票, 利润为 0
• $f(-1,k,0)$: 不在股票交易日, 持有0份股票, 利润为 0
• $f(i,0,1)$: 最大交易次数为0, 持有1份股票, 没有意义, 设定为 -Infinity
• $f(-1,k,1)$: 不在股票交易日, 持有1份股票, 没有意义, 设定为 -Infinity

得出了通用状态转移方程后, 就可以解决不同情况下的股票问题了.

1.1 买卖股票的最佳时机 (K=1)

每天只能进行一次操作, 即k=1. 而得出状态转移方程如下:

$\begin{cases} f(i,1,0)=max(f(i-1,1,0),f(i-1,1,1)+Prices[i]) \\ f(i,1,1)=max(f(i-1,1,1),f(i-1,1-1,0)-Prices[i])=max(f(i-1,1,1),-Prices[i]) \end{cases}$

1.1.1 迭代+二维数组

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    // DP
    dp := make([][2]int, len(prices))

    dp[0][0] = 0
    dp[0][1] = -prices[0]

    for i := 1; i < len(prices); i++ {
        dp[i][0] = max(dp[i-1][0], dp[i-1][1]+prices[i])
        dp[i][1] = max(dp[i-1][1], -prices[i])
    }

    return dp[len(prices)-1][0]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.1.2 迭代+一维数组(两个辅助变量)

因为每次计算需要左上方dp[i-1][j-1]和正上方结果dp[i-1][j], 可以优化到一维数组, 再进一步优化到直接使用两个变量即可.

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    n := len(prices)
    // DP
    dp := make([]int, 2)
    dp[0] = 0
    dp[1] = -prices[0]

    for i := 1; i < n; i++ {
        dp[0] = max(dp[0], dp[1]+prices[i])
        dp[1] = max(dp[1], -prices[i])
    }

    return dp[0]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

两个辅助变量

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    // DP
    profit0 := 0
    profit1 := -prices[0]

    for i := 1; i < len(prices); i++ {
        profit0 = max(profit0, profit1 + prices[i])
        profit1 = max(profit1, -prices[i])
    }

    return profit0
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.2 买卖股票的最佳时机 II (K=+Infinity)

K 为正无穷, 此时k和k-1视作相同. 那么状态转移方程为:

$\begin{cases} f(i,k,0)=max(f(i-1,k,0),f(i-1,k,1)+Prices[i]) \\ f(i,k,1)=max(f(i-1,k,1),f(i-1,k-1,0)-Prices[i]) \end{cases}$

1.2.1 迭代+二维数组

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    // DP
    dp := make([][2]int, len(prices))
    dp[0][0] = 0
    dp[0][1] = -prices[0]

    for i := 1; i < len(prices); i++ {
        dp[i][0] = max(dp[i-1][0], dp[i-1][1]+prices[i])
        dp[i][1] = max(dp[i-1][1], dp[i-1][0]-prices[i])
    }

    return dp[len(prices)-1][0]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.2.2 迭代+一维数组(多个辅助变量)

计算结果只需要前一个数据, 优化到一维数组.

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    // DP
    dp := make([]int, 2)
    dp[0] = 0
    dp[1] = -prices[0]

    prev := dp[0]
    for i := 1; i < len(prices); i++ {
        cur := max(dp[0], dp[1]+prices[i])
        dp[1] = max(dp[1], prev-prices[i])
        prev = cur
        dp[0] = cur
    }

    return dp[0]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

优化到多个辅助变量

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    // DP
    profit0 := 0
    profit1 := -prices[0]

    prev := profit0
    for i := 1; i < len(prices); i++ {
        cur := max(profit0, profit1+prices[i])
        profit1 = max(profit1, prev-prices[i])
        prev = cur
        profit0 = cur
    }

    return profit0
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

这种写法要更加直观

   for i := 1; i < len(prices); i++ {
        newProfit0 := max(profit0, profit1+prices[i])
        newProfit1 := max(profit1, profit0-prices[i])
        profit0 = newProfit0
        profit1 = newProfit1
    }

1.3 买卖股票的最佳时机 III (K=2)

此时状态转移方程有四组:

$\begin{cases} f(i,2,0)=max(f(i-1,2,0), f(i-1,2,1)+Prices[i])\\ f(i,2,1)=max(f(i-1,2,1),f(i-1,1,0)-Prices[i])\\ f(i,1,0)=max(f(i-1,1,0),f(i-1,1,1)+Prices[i]) \\ f(i,1,1)=max(f(i-1,1,1),-Prices[i]) \end{cases}$

1.3.1 迭代+三维数组

根据状态转移公式得出代码(三维数组, 🧠开始冒烟了):

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    n := len(prices)
    // DP
    dp := make([][3][2]int, n)
    
    // 初始值
    dp[0][2][0] = 0
    dp[0][2][1] = -prices[0]
    dp[0][1][0] = 0
    dp[0][1][1] = -prices[0]

    // 计算
    for i := 1; i < n; i++ {
        dp[i][2][0] = max(dp[i-1][2][0], dp[i-1][2][1]+prices[i])
        dp[i][2][1] = max(dp[i-1][2][1], dp[i-1][1][0]-prices[i])
        dp[i][1][0] = max(dp[i-1][1][0], dp[i-1][1][1]+prices[i])
        dp[i][1][1] = max(dp[i-1][1][1], -prices[i])
    }

    return dp[n-1][2][0]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.3.2 迭代+一维

因为计算还是只需要前一个数据, 那么可以优化到多个变量.

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    n := len(prices)
    // DP
    // 初始值
    profitTwo0 := 0
    profitTwo1 := -prices[0]
    profitOne0 := 0
    profitOne1 := -prices[0]

    // 计算
    for i := 1; i < n; i++ {
        profitTwo0 = max(profitTwo0, profitTwo1+prices[i])
        profitTwo1 = max(profitTwo1, profitOne0-prices[i])
        profitOne0 = max(profitOne0, profitOne1+prices[i])
        profitOne1 = max(profitOne1, -prices[i])
    }

    return profitTwo0
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.4 买卖股票的最佳时机 IV (K 可以是任意值)

这个最普遍的情况, 因为一次交易最短为两天(一天买入, 一天卖出), 那么有利润的最大操作数为n/2.

• 当k >= n/2时, 此时最大收益不由k决定了, k和k等于正无穷时是一样的.

• 当k < n/2 时, 此时使用状态转移公式计算最终结果$f(n-1,k,0)$

1.4.1 迭代+三维数组 TC: O(nk), SC: O(nk)

func maxProfit(k int, prices []int) int {
    if len(prices) == 0 || k == 0 {
        return 0
    }
    n := len(prices)
    
    if k >= n/2 {
        return maxInfTransProfit(prices)
    }
    
    // DP
   dp := make([][][2]int, n) 
    for i := range dp {
        dp[i] = make([][2]int, k+1)
    }
    // 每种k的初始值
    for i := 1; i <= k; i++ {
        dp[0][i][0] = 0
        dp[0][i][1] = -prices[0]
    }

    for i := 1; i < n; i++ {
        for j := 1; j <= k; j++ {
            dp[i][j][0] = max(dp[i-1][j][0], dp[i-1][j][1] + prices[i])
            dp[i][j][1] = max(dp[i-1][j][1], dp[i-1][j-1][0] - prices[i])
        }
    }

    return dp[n-1][k][0]
}

func maxInfTransProfit(prices []int) int {
    // DP
    profit0 := 0
    profit1 := -prices[0]

    for i := 1; i < len(prices); i++ {
        profit1 = max(profit1, profit0-prices[i])
        profit0 = max(profit0, profit1+prices[i])
    }

    return profit0
}


func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.4..2 迭代+二维数组+反向计算 SC: O(k)

计算只需左边的数据, 可以反向计算.

此时优化为二维数组,

func maxProfit(k int, prices []int) int {
    if len(prices) == 0 || k == 0 {
        return 0
    }
    n := len(prices)
    
    if k >= n/2 {
        return maxInfTransProfit(prices)
    }
    
    // DP
    dp := make([][2]int, k+1)
    // 每种k的初始值
    for i := 1; i <= k; i++ {
        dp[i][0] = 0
        dp[i][1] = -prices[0]
    }

    for i := 1; i < n; i++ {
        for j := k; j >= 1; j-- {
            dp[j][0] = max(dp[j][0], dp[j][1]+prices[i])
            dp[j][1] = max(dp[j][1], dp[j-1][0]-prices[i])
        }
    }

    return dp[k][0]
}

func maxInfTransProfit(prices []int) int {
    // DP
    profit0 := 0
    profit1 := -prices[0]

    for i := 1; i < len(prices); i++ {
        profit1 = max(profit1, profit0-prices[i])
        profit0 = max(profit0, profit1+prices[i])
    }

    return profit0
}


func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.5 最佳买卖股票时机含冷冻期(1天) (K=+Infinity)

因为包含冷冻期(1天), 那么在i-1天进行交易后, 第i天不能交易. 此时状态转移方程如下:

$\begin{cases} f(i,k,0)=max(f(i-1,k,0),f(i-1,k,1)+Prices[i]) \\ f(i,k,1)=max(f(i-1,k,1),f(i-2,k-1,0)-Prices[i]) \end{cases}$

i-1变成了i-2

1.5.1 迭代+二维数组

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    n := len(prices)
    // DP
    dp := make([][2]int, n)
    dp[0][0] = 0
    dp[0][1] = -prices[0]

    for i := 1; i < n; i++ {
        dp[i][0] = max(dp[i-1][0], dp[i-1][1]+prices[i])
        if i >= 2 {
            dp[i][1] = max(dp[i-1][1], dp[i-2][0]-prices[i])
        } else {
            dp[i][1] = max(dp[i-1][1], -prices[i])
        }
    }

    return dp[n-1][0]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.5.2 迭代+一维

计算i-1和i-2有关, 使用辅助变量暂存

func maxProfit(prices []int) int {
    if len(prices) == 0 {
        return 0
    }
    n := len(prices)
    // DP
    profit0 := 0
    prevProfit0 := 0
    profit1 := -prices[0]
    for i := 1; i < n; i++ {
       nextProfit0 := max(profit0, profit1+prices[i])
       if i >= 2 {
           profit1 = max(profit1, prevProfit0-prices[i])
       } else {
           profit1 = max(profit1, -prices[i])
       }
       prevProfit0 = profit0
       profit0 = nextProfit0
    }

    return profit0
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.6 买卖股票的最佳时机含手续费 (K=+Infinity)

每次买入, 卖出都需要减去手续费, 修改状态转移方程如下:

$\begin{cases} f(i,k,0)=max(f(i-1,k,0),f(i-1,k,1)+Prices[i]) \\ f(i,k,1)=max(f(i-1,k,1),f(i-1,k-1,0)-Prices[i]-fee) \end{cases}$

1.6.1迭代+二维数组

func maxProfit(prices []int, fee int) int {
    if len(prices) == 0 {
        return 0
    }
    n := len(prices)
    // DP
    dp := make([][2]int, n)
    dp[0][0] = 0
    dp[0][1] = -prices[0]-fee

    for i := 1; i < n; i++ {
        dp[i][0] = max(dp[i-1][0], dp[i-1][1]+prices[i])
        dp[i][1] = max(dp[i-1][1], dp[i-1][0]-prices[i]-fee)
    }

    return dp[n-1][0]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

1.6.2 迭代+一维

func maxProfit(prices []int, fee int) int {
    if len(prices) == 0 {
        return 0
    }
    n := len(prices)
   
    profit0 := 0
    profit1 := -prices[0]-fee
    for i := 1; i < n; i++ {
       newProfit0 := max(profit0, profit1+prices[i])
       newProfit1 := max(profit1, profit0-prices[i]-fee)
       profit0 = newProfit0
       profit1 = newProfit1 
    }

    return profit0
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

2. 假设 卖出 操作消耗最大交易次数

状态转移方程和卖出时有少许差别:

$\begin{cases} f(i,k,0)=max(f(i-1,k,0), f(i-1,k-1,1)+Prices[i]) \\ f(i,k,1)=max(f(i-1,k,1), f(i-1,k,0)-Prices[i]) \end{cases}$

只有在卖出时, 最大操作数才被消耗.

Reference

1. https://leetcode.com/problems/best-time-to-buy-and-sell-stock-with-transaction-fee/solutions/108870/Most-consistent-ways-of-dealing-with-the-series-of-stock-problems/
2. https://leetcode.com/problems/best-time-to-buy-and-sell-stock-with-transaction-fee/solutions/108870/Most-consistent-ways-of-dealing-with-the-series-of-stock-problems/
3. https://leetcode.cn/circle/discuss/qiAgHn/