0309. Best Time to Buy and Sell Stock with Cooldown

Medium | Array + DP | 32 ms (96.27%), 14.6 MB (38.63%)

❓ Problem Statement

Source: LeetCode - Best Time to Buy and Sell Stock with Cooldown GitHub: Solution / Performance

You are given an array prices where prices[i] is the price of a given stock on the ith day.

Find the maximum profit you can achieve. You may complete as many transactions as you like (i.e., buy one and sell one share of the stock multiple times) with the following restrictions:

• After you sell your stock, you cannot buy stock on the next day (i.e., cooldown one day).

Note: You may not engage in multiple transactions simultaneously (i.e., you must sell the stock before you buy again).

✍🏻 Constraints & Example

Constraints:

• 1 <= prices.length <= 5000

• 0 <= prices[i] <= 1000

Input: prices = [1,2,3,0,2]
Output: 3
Explanation: transactions = [buy, sell, cooldown, buy, sell]

Input: prices = [1]
Output: 0

💡 Ideas

We can use Final State Machine to solve the series of these problems.

Note that this problem limits us to have one day cooldown after selling a stock (belongs to no-hold state), so for calculating of the hold state's profit, we need the other variable to record the profit whenever we sell a stock.

This variable could be initialized as 0 (to simulate as we buy and sell stock in the same day long time ago). Then, this variable is updated by previous no-hold state's profit to create a one-day gap.

• Hold state means you hold the stock.

  • Since you don't have any stocks at the beginning, the profit at this state is initialized with negative infinity.

  • The profit in this state is calculated by: previous no-hold state's profit - current stock price (= buy) = previous profit when we sell stock - current stock price

• No-Hold state means you do not hold any stocks.

  • Since you don't have any stocks at the beginning (you're initially at this state ), the profit at this state is initialized with 0.

  • The profit in this state is calculated by: previous hold state's profit + current stock price (= sell)

  • Return the profit at the no-hold state as the maximum profit since the hold state still has stock on hold.

🤖 Python3

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        #  (base case)
        if len(prices) == 0 or len(prices) == 1: return 0
        
        # ==================================================
        #  Array + Dynamic Programming              (FSM)  =
        # ==================================================
        # time  : O(n)
        # space : O(1)

        hold, noHold = float('-inf'), 0
        
        #  after selling stock, you cannot buy stock on the next day = NO HOLD state
        preDaySell = 0
        
        for price in prices:
            preHold, preNoHold = hold, noHold
            
            #  due to cooldown, cannot use preNoHold's profit to buy stock, use preDaySell
            hold   = max(preHold,   preDaySell - price)
            noHold = max(preNoHold, preHold    + price)
            
            #  record at the end to have an one-day gap (still 0 after 1st iteration)
            preDaySell = preNoHold
            
        return noHold

🤖 Java

class Solution {
    /**
     * @time  : O(n)
     * @space : O(1)
     */
     
    public int maxProfit(int[] prices) {
        /* base case */
        if(prices.length == 1) return 0;
        
        int hold = Integer.MIN_VALUE, noHold = 0;
        int preDaySell = 0;
        
        for(int price: prices) {
            int preHold = hold, preNoHold = noHold;
            
            hold   = Math.max(preHold,   preDaySell - price);
            noHold = Math.max(preNoHold, preHold    + price );
            
            preDaySell = preNoHold;
        }
        
        return noHold;
    }
}