合縱連橫: 從DP框架理解 最佳股票買賣系列題 的背後本質

Obi - @pixel8propix on Unsplash

這篇文章，會帶著大家複習以前學過的FSM+DP框架，
並且以有限狀態機 + DP狀態轉移的概念為核心，

貫穿一些相關聯最佳股票買賣系列的題目，
透過框架複現來幫助讀者理解這個實用的演算法框架。

基本的FSM + DP 框架，配合交易邏輯。

針對每一天，其實歸根究柢只有兩種狀態。

第一種是手上都持有現金。

第二種是手上持有股票。

每天進行的動作，也只有兩種可能情況:

假如今天結束的時候，是都持有現金。
代表昨天就已經持有現金今天繼續持有；或者，昨天持有股票，今天選擇賣出股票。

假如今天結束的時候，是持有股票。
代表昨天就已經持有股票今天繼續抱著；或者，昨天持有現金，今天選擇買入股票。

寫成虛擬碼或者演算法的話，可以這樣表達

DP[day_d, KeepCash] 
= max( DP[day_d-1, KeepCash], DP[day_d-1, HoldStack] + stockPrice[day_d] )

DP[day_d, HoldStock]
= max( DP[day_d-1, HoldStock], DP[day_d-1, KeepCash] - stockPrice[day_d] )

接下來，我們會用這個上面這種結構，貫穿一些同類型，有關聯的題目

(請讀者、或觀眾留意上面的這種 FSM+ DP 動態規劃 框架，之後的解說會反覆出現)

FSM有限狀態機 + DP 動態規劃框架 示意圖

打鐵趁熱，現在來看最佳股票買賣系列題第一題。

Best time to buy and sell stock

這題的限制條件是只能做一組交易(也就是一次買入，一次賣出)。

只有一次買入，相當於會從某一天選擇買入，之後就只能選另一天賣出了。

某一天買入之前一定是沒賺沒賠，所以寫成式子就是 0 - stock_price

最後用FSM+DP框架的思考邏輯寫成程式碼，就是下方這個樣子

    
class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        
        # It is impossible to have stock to sell on first day, so -infinity is set as initial value
        cur_hold, cur_cash = -float('inf'), 0
        
        
        for stock_price in prices:
            
            prev_hold, prev_cash = cur_hold, cur_cash
            
            # either keep in hold, or just buy today with stock price
            cur_hold = max(prev_hold, 0-stock_price)
            
            # either keep in cash, or just sell today with stock price
            cur_cash = max(prev_cash, prev_hold + stock_price)
            
            
        # max profit must be in Cash state
        return cur_cash 

接著來看變化題，最佳股票買賣系列題第二題。

Best time to buy and sell stock II

這題比較寬鬆，不限制交易組數，也就是可以交易任意多組，
買入股票n次，賣出股票n次。

也就是說，某一天才買入的狀態會相依於前一天現金部位的狀態

所以寫成式子就是 dp[day-1, KEEP_CASH] - stock_price

最後用FSM+DP框架的思考邏輯寫成程式碼，就是下方這個樣子

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        
        # Use constant literal to help reader understand source code below.
        HOLD_STOCK, KEEP_CASH = 0, 1

        ## Dictionary (Hash table)
        # key: (day, state) pair
        # value: coreesponding maximal profit
        dp = {}

        # No free lunch, it is impoosible to have stock before first trading day
        dp[-1, HOLD_STOCK] = -math.inf

        # No gain, no loss before first trading day
        dp[-1, KEEP_CASH] = 0

        # For each day with corresponding stock price in stock market
        for day, stock_price in enumerate(prices):

            # If today we have stock, either we already had it yesterday, or we just buy stock and hold it today.
            dp[day, HOLD_STOCK] = max( dp[day-1, HOLD_STOCK], dp[day-1, KEEP_CASH] - stock_price)

            # If today we keep cash, either we already kept cash yesterday, or we just sell out stock today
            dp[day, KEEP_CASH] = max( dp[day-1, KEEP_CASH], dp[day-1, HOLD_STOCK] + stock_price)
        
        #---------------------------------------------------------------
        # To get maximal realized profit, final state must be KEEP_CASH.
        last_day = len(prices)-1
        return dp[last_day, KEEP_CASH]

Best time to buy and sell stock II 空間優化的程式碼實作

仔細觀察，可以發現每天的狀態只依賴前一天的狀態去更新。

因此，可以用之前學過的狀態壓縮的技巧，把表格拿掉，
改成只用變數儲存狀態，節省空間，從O(n)空間降到O(1)常數級別。

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        
					# It is impossible to sell stock on first day, set -infinity as initial value for cur_hold
        cur_hold, cur_cash = -float('inf'), 0
        
        for stock_price in prices:
            
            prev_hold, prev_cash = cur_hold, cur_cash
            
							# either keep hold, or buy in stock today at stock price
            cur_hold = max( prev_hold, prev_cash - stock_price )
			
							# either keep cash, or sell out stock today at stock price
            cur_cash = max( prev_cash, prev_hold + stock_price )
            
        # maximum profit must be in Cash state
        return cur_cash

有了狀態圖以後，再回來看程式碼清晰許多了，對吧!

接下來，為了再更逼真一點，延伸題就是把交易的額外成本fee(稅, 手續費)也加入考量。

延伸題就是

Best time to buy and sell stock with transaction fee

那也很容易，我們就在每組交易進行的時候，額外多扣一個費用fee即可。

註: 不失一般性，我們選擇在買入股票時付出額外的交易成本費用fee。

讀者想要在賣出時扣掉也可以，可以試著改改看，也能通過測試！

Best time to buy and sell stock with transaction fee 程式碼實作

class Solution:
    def maxProfit(self, prices: List[int], fee: int) -> int:
        
        cur_hold, cur_cash = -math.inf, 0
        
        for stock_price in prices:
            
            prev_hold, prev_cash = cur_hold, cur_cash
            
            # either keep cash, or sell out today at stock price
            cur_cash = max(prev_cash, prev_hold + stock_price )
            
            # either keep hold, or buy in today at stock price and pay transaction fee for this trade
            cur_hold = max(prev_hold, prev_cash - stock_price - fee)
        
        # maximum profit must be in not-hold state
        return cur_cash

好，接著來看進階題，最佳股票買賣系列題第三題。

Best time to buy and sell stock III

這題的限制條件是最多只能兩組交易(也就是各兩次買入，各兩次賣出)。

也就是說，第i次的買入並且持有股票的狀態，
前面會相依於第(i-1)次賣出並且持有現金的狀態。

所以寫成式子就是 dp[KEEP_CASH, i-1] - stock_price

最後用FSM+DP框架的思考邏輯寫成程式碼，就是下方這個樣子

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        
        HOLD_STOCK, KEEP_CASH = 0, 1
        
        # dictionary
        # key: state, kth-transaction
        # value: corresponding profit
        dp = defaultdict(int)        
        
        # No free lunch, impoosible to have stock before first trading day
        dp[(HOLD_STOCK, 0)] = -math.inf
        dp[(HOLD_STOCK, 1)] = -math.inf
        dp[(HOLD_STOCK, 2)] = -math.inf
        

        for stock_price in prices:

            ## For 1st transaction:
            # Either we kept cash already, or we just sell out stock today
            dp[KEEP_CASH, 1] = max(dp[KEEP_CASH, 1], dp[HOLD_STOCK, 1] + stock_price )
            
            # Either we had stock already, or we just buy in stock today ( add one more transaction)
            dp[HOLD_STOCK, 1] = max(dp[HOLD_STOCK, 1],  dp[KEEP_CASH, 0] - stock_price)
            
            ## For 2nd transaction:
            # Either we kept cash already, or we just sell out stock today
            dp[KEEP_CASH, 2] = max(dp[KEEP_CASH, 2], dp[ HOLD_STOCK, 2] + stock_price)

            # Either we had stock already, or we just buy in stock today ( add one more transaction)
            dp[HOLD_STOCK, 2] = max(dp[HOLD_STOCK, 2], dp[KEEP_CASH, 1] - stock_price)
            
        # Maximal profit must be KEEP_CASH on last day
        # (This means we cash out and sell stocks finally)
        return dp[KEEP_CASH, 2]

仔細觀察，其實for loop裡面包含著一個階梯迭代更新的關係。

第i組的交易狀態更新，只依賴第i組和第i-1組的狀態。

抽象化+結構化，可以寫成這個形式

for trans in range(1, k+1):

	dp[HOLD_STOCK, trans] = max( dp[HOLD_STOCK, trans], dp[KEEP_CASH, trans-1 ] - stock_price )

	dp[KEEP_CASH, trans] = max( dp[KEEP_CASH, trans], dp[HOLD_STOCK, trans] + stock_price )

寫成這個形式，也相當於推廣成

最多只能k組交易下的最佳股票買賣模式

當我們推導出這個通則的時候，其實無形中已經把下一題也解出來了！ XD

下一題就是

Best time to buy and sell stock IV

這題的限制條件是最多只能k組交易(也就是各k次買入，各k次賣出)。

這邊承接第剛剛上面推導出來的通則，順著寫下來即可。

class Solution:
    def maxProfit(self, k: int, prices: List[int]) -> int:
        
        n = len(prices)
        
        HOLD_STOCK, KEEP_CASH = 0, 1
        
        dp = defaultdict(int)
        for k in range (0, k+1):
            dp[HOLD_STOCK, k] = -math.inf

        for stock_price in prices:
            
            for trans in range(1, k+1):
                dp[HOLD_STOCK, trans] = max( dp[HOLD_STOCK, trans], dp[KEEP_CASH, trans-1 ] - stock_price )
                
                dp[KEEP_CASH, trans] = max( dp[KEEP_CASH, trans], dp[HOLD_STOCK, trans] + stock_price )

        return dp[KEEP_CASH, k]

萬變不離其宗，其實DP狀態轉移的觀念其實都是相通的。

最後一題是一個延伸變化題

Best Time to Buy and Sell Stock with Cooldown

這題的也是不限制交易組數，也就是可以交易任意多組，

買入股票n次，賣出股票n次。

也就是說，某一天才買入的狀態會相依於前一天現金部位的狀態

但是，多了一個冷卻條件，
股票賣出之後，會強迫進入冷卻日一天，冷卻日的次日才可以再買入股票。

畫成FSM有限狀態機的狀態圖就是

順著狀態圖的更新邏輯寫下迭代耕新的程式碼即可

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        
				# initialization
        cool_down, sell, hold = 0, 0, -float('inf')
        
        for stock_price_of_Day_i in prices:
            
            prev_cool_down, prev_sell, prev_hold = cool_down, sell, hold
            
            # Max profit of cooldown on Day i comes from either cool down of Day_i-1, or sell out of Day_i-1 and today Day_i is cooling day
            cool_down = max(prev_cool_down, prev_sell)
            
            # Max profit of sell on Day_i comes from hold of Day_i-1 and sell on Day_i
            sell = prev_hold + stock_price_of_Day_i
            
            # Max profit of hold on Day_i comes from either hold of Day_i-1, or cool down on Day_i-1 and buy on Day_i
            hold = max(prev_hold, prev_cool_down - stock_price_of_Day_i)
        
        
        # The action of final trading day must be either sell or cool down
        return max(sell, cool_down)

可以視為Best time to buy and sell stock II的延伸題(多了額外的冷卻條件)

結語

好，今天一口氣介紹了最精華的部分，

通用的FSM有限狀態機 + DP狀態轉移框架，還有相關的衍伸變化題與演算法建造流程。

希望能幫助讀者、觀眾徹底理解它的原理，

並且能夠舉一反三，洞察其背後的本質和了解相關的應用領域！

感謝收看囉! 我們下篇文章再見~