One Method to Solve all Stock Problems on LeetCode

Many readers complain that there are too many solutions to the stock series problems on LeetCode. If you encounter such problems in an interview, you probably won't think of those clever methods. So what should you do? Therefore, this article does not focus on overly clever approaches but instead adopts a solid and steady method to solve all problems with a universal approach, meeting changes with constancy.

This article refers to the ideas from the highly-rated English solution and uses the state machine technique to solve the problems, which can pass all submissions. Don't be intimidated by this term; it's just a literary expression. In reality, it's just a DP table, and you'll understand it at a glance.

Let's randomly pick a problem and look at other people's solutions:

int maxProfit(int[] prices) {
    if(prices.empty()) return 0;
    int s1 = -prices[0], s2 = INT_MIN, s3 = INT_MIN, s4 = INT_MIN;

    for(int i = 1; i < prices.size(); ++i) {            
        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);
}

Can you understand it? Can you do it? No way, you can't understand it, and that's normal. Even if you barely understand it, you still won't be able to solve the next problem. Why can others come up with such strange yet efficient solutions? Because these problems follow a framework, but no one will tell you about it. Once they do, you'll learn it in five minutes, and the mystery of these algorithm questions will vanish, making them easily solvable.

This article will reveal that framework to you and guide you through solving each problem effortlessly. We'll use the technique of state machines to solve these problems, and all solutions will pass the tests. Don't be intimidated by the term "state machine"; it's just a fancy word for a DP table, and you'll understand it at a glance.

These 6 problems share common characteristics. We only need to focus on LeetCode problem 188, "Best Time to Buy and Sell Stock IV," because it is the most generalized form. Other problems are simplified versions of this one. Let's look at the problem:

Input: k = 2, prices = [2,4,1]
Output: 2
Explanation: Buy on day 1 (price = 2) and sell on day 2 (price = 4), profit = 4-2 = 2.

Input: k = 2, prices = [3,2,6,5,0,3]
Output: 7
Explanation: Buy on day 2 (price = 2) and sell on day 3 (price = 6), profit = 6-2 = 4. Then buy on day 5 (price = 0) and sell on day 6 (price = 3), profit = 3-0 = 3.

The first problem involves only one transaction, equivalent to k = 1; the second problem allows unlimited transactions, equivalent to k = +infinity (positive infinity); the third problem involves two transactions, equivalent to k = 2; the remaining two problems also allow unlimited transactions but add extra conditions like a "cooldown period" and "transaction fees," which are just variations of the second problem and easy to handle.

Now, let's get back to solving the problems.

Chapter 1: Exhaustive Framework

First, let's tackle the same question: How to exhaust all possibilities?

As mentioned in Dynamic Programming Core Techniques, dynamic programming essentially involves exhaustively listing all possible "states" and then selecting the optimal solution from the "choices".

For this problem, let's break it down day by day to see how many possible "states" there are and identify the corresponding "choices" for each state. Our goal is to exhaust all "states" to update them based on the corresponding "choices". It might sound abstract, but just remember the terms "state" and "choice". It will become clear with a practical example.

for state1 in all_values_of_state1:
    for state2 in all_values_of_state2:
        for ...
            dp[state1][state2][...] = choose_the_best(choice1, choice2...)

In this problem, there are three "choices" each day: buy, sell, and do nothing (rest). We use buy, sell, and rest to represent these choices.

However, not all choices are available every day. For instance, sell must follow buy, and buy must follow sell. Therefore, the rest action should also have two states: one after buy (holding a stock) and one after sell (not holding a stock). Don't forget, we also have a limit on the number of transactions k, meaning you can only buy if k > 0.

Sounds complicated, right? Don't worry. Our goal right now is just to enumerate everything. No matter how many states you have, my job is to list them all out in one go.

This problem has three "states": the first is the number of days, the second is the maximum number of transactions allowed, and the third is the current holding state (i.e., the previously mentioned rest state, which we can represent as 1 for holding and 0 for not holding). We can use a three-dimensional array to store all combinations of these states:

dp[i][k][0 or 1]
0 <= i <= n - 1, 1 <= k <= K
n is the number of days, K is the upper limit of transactions, and 0 and 1 represent whether you hold a stock.
This problem has n × K × 2 states in total, and we can solve it by enumerating all of them.

for 0 <= i < n:
    for 1 <= k <= K:
        for s in {0, 1}:
            dp[i][k][s] = max(buy, sell, rest)

Moreover, we can describe the meaning of each state in natural language. For example, dp[3][2][1] means: it's the third day, I currently hold a stock, and I have made at most 2 transactions so far. Another example, dp[2][3][0] means: it's the second day, I currently do not hold any stock, and I have made at most 3 transactions so far. It's quite easy to understand, right?

The final answer we want is dp[n - 1][K][0], which means on the last day, with at most K transactions allowed, the maximum profit we can achieve.

You might wonder why not dp[n - 1][K][1]? Because dp[n - 1][K][1] means holding a stock on the last day, while dp[n - 1][K][0] means the stock has been sold by the last day. Obviously, the latter will yield a higher profit than the former.

Remember how to interpret "states". Whenever you find something hard to understand, translating it into natural language can make it easier to grasp.

2. State Transition Framework

Now that we have completed the enumeration of "states," we begin to consider the "choices" available in each "state" and how to update the "state."

Focusing only on the "holding state," we can draw a state transition diagram:

From this diagram, it is clear how each state (0 and 1) transitions. Based on this, we can write the state transition equations:

dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i])
              # max( Choose to rest today, Choose to sell today )

Explanation: Today, I do not hold any stock. There are two possibilities from which I seek the maximum profit:

1. I did not hold any stock yesterday and the maximum transaction limit up to yesterday was k; today I choose to rest, so I still do not hold any stock and the transaction limit remains k.

2. I held stock yesterday and the transaction limit up to yesterday was k; but today I sell, so I do not hold any stock today, and the transaction limit remains k.

dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i])
              # max( Choose to rest today, Choose to buy today )

Explanation: Today, I hold stock with a transaction limit of k. For yesterday, there are two possibilities from which I seek the maximum profit:

1. I held stock yesterday and the transaction limit up to yesterday was k; today I choose to rest, so I still hold the stock, and the transaction limit remains k.

2. I did not hold stock yesterday and the transaction limit up to yesterday was k - 1; but today I choose to buy, so I hold stock today, and the transaction limit is k.

This explanation should be clear: when you buy, you subtract prices[i] from the profit; when you sell, you add prices[i] to the profit. Today's maximum profit is the larger of these two choices.

Note the restriction of k: choosing to buy initiates a transaction, so yesterday's transaction limit k should decrease by 1.

Now, we have completed the most challenging step in dynamic programming: the state transition equation. If you understood the previous content, you can already solve all problems quickly by applying this framework. However, there's one last bit missing, which is defining the base case, the simplest scenario.

dp[-1][...][0] = 0
Explanation: Since `i` starts from 0, `i = -1` means we haven't started yet, so the profit is naturally 0.

dp[-1][...][1] = -infinity
Explanation: Before we start, it's impossible to hold any stock.
Since our algorithm requires a maximum value, we set the initial value to the smallest possible value to easily find the maximum.

dp[...][0][0] = 0
Explanation: Since `k` starts from 1, `k = 0` means no transactions are allowed, so the profit is naturally 0.

dp[...][0][1] = -infinity
Explanation: When no transactions are allowed, it's impossible to hold any stock.
Again, we set the initial value to the smallest possible value for ease of finding the maximum.

Summarizing the state transition equations:

Base case:
dp[-1][...][0] = dp[...][0][0] = 0
dp[-1][...][1] = dp[...][0][1] = -infinity

State transition equations:
dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i])
dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i])

You might be wondering how to represent an array index of -1 in programming, or how to represent negative infinity. These are detail issues and there are many ways to handle them. Now that the complete framework is in place, let's move on to specific examples.

Three, Quick Win Problems

121. Best Time to Buy and Sell Stock

First problem, let's talk about LeetCode problem 121, 'Best Time to Buy and Sell Stock,' which is equivalent to the case when k = 1:

Input: prices = [7,1,5,3,6,4]
Output: 5
Explanation: Buy on day 2 (price = 1) and sell on day 5 (price = 6), profit = 6-1 = 5.
Note that buying on day 2 and selling on day 1 is not allowed because you must buy before you sell.

Input: prices = [7,6,4,3,1]
Output: 0
Explanation: In this case, no transactions are done and the max profit = 0.

Directly apply the state transition equation and simplify based on the base case:

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], dp[i-1][0][0] - prices[i]) 
            = max(dp[i-1][1][1], -prices[i])
Explanation: For the base case when k = 0, dp[i-1][0][0] = 0.

Now we notice that k is always 1, which means it doesn't change and no longer affects the state transition.
We can further simplify by removing all k:
dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
dp[i][1] = max(dp[i-1][1], -prices[i])

Here is the direct code implementation:

int n = prices.length;
int[][] dp = new int[n][2];
for (int i = 0; i < n; i++) {
    dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
    dp[i][1] = Math.max(dp[i-1][1], -prices[i]);
}
return dp[n - 1][0];

int n = prices.size();
vector<vector<int>> dp(n, vector<int>(2));
dp[0][0] = 0;
dp[0][1] = -prices[0];

for (int 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], -prices[i]);
}

return dp[n-1][0];

n = len(prices)
dp = [[0]*2 for _ in range(n)]
for i in range(n):
    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[n - 1][0]

n := len(prices)
dp := make([][2]int, n)
for i := 0; i < n; i++ {
    if i == 0 {
        dp[i][0] = 0
        dp[i][1] = -prices[i]
        continue
    }
    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[n-1][0]

var n = prices.length;
var dp = Array.from({ length: n }, () => new Array(2).fill(0));
for (var i = 0; i < n; i++) {
    dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
    dp[i][1] = Math.max(dp[i-1][1], -prices[i]);
}
return dp[n - 1][0];

Clearly, when i = 0, i - 1 is an invalid index. This is because we haven't handled the base case for i. We can give a special treatment like this:

if (i - 1 == -1) {
    dp[i][0] = 0;
    // According to the state transition equation:
    //   dp[i][0] 
    // = max(dp[-1][0], dp[-1][1] + prices[i])
    // = max(0, -infinity + prices[i]) = 0
    //   dp[i][0] 
    // = max(dp[-1][0], dp[-1][1] + prices[i])
    // = max(0, -infinity + prices[i]) = 0

    dp[i][1] = -prices[i];
    // According to the state transition equation:
    //   dp[i][1] 
    // = max(dp[-1][1], dp[-1][0] - prices[i])
    // = max(-infinity, 0 - prices[i]) 
    // = -prices[i]
    //   dp[i][1] 
    // = max(dp[-1][1], dp[-1][0] - prices[i])
    // = max(-infinity, 0 - prices[i]) 
    // = -prices[i]
    continue;
}

The first problem is solved, but handling the base case this way is quite troublesome. Also, notice the state transition equation: the new state only depends on the adjacent state. Therefore, we can use the technique from the previous article Dimensionality Reduction in Dynamic Programming: Space Compression Techniques. Instead of using the entire dp array, we only need a single variable to store the adjacent state. This reduces the space complexity to O(1):

// Original version
int maxProfit_k_1(int[] prices) {
    int n = prices.length;
    int[][] dp = new int[n][2];
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i-1][1], -prices[i]);
    }
    return dp[n - 1][0];
}

// Space complexity optimized version
int maxProfit_k_1(int[] prices) {
    int n = prices.length;
    // base case: dp[-1][0] = 0, dp[-1][1] = -infinity
    int dp_i_0 = 0, dp_i_1 = Integer.MIN_VALUE;
    for (int i = 0; i < n; i++) {
        // dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        // dp[i][1] = max(dp[i-1][1], -prices[i])
        dp_i_1 = Math.max(dp_i_1, -prices[i]);
    }
    return dp_i_0;
}

// Original version
int maxProfit_k_1(vector<int>& prices) {
    int n = prices.size();
    vector<vector<int>> dp(n, vector<int>(2));
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        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[n - 1][0];
}

// Space complexity optimized version
int maxProfit_k_1(vector<int>& prices) {
    int n = prices.size();
    // base case: dp[-1][0] = 0, dp[-1][1] = -infinity
    int dp_i_0 = 0, dp_i_1 = INT_MIN;
    for (int i = 0; i < n; i++) {
        // dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i]);
        // dp[i][1] = max(dp[i-1][1], -prices[i])
        dp_i_1 = max(dp_i_1, -prices[i]);
    }
    return dp_i_0;
}

# Original version
def maxProfit_k_1(prices: List[int]) -> int:
    n = len(prices)
    dp = [[0 for i in range(2)] for j in range(n)]
    for i in range(n):
        if i - 1 == -1:
            # base case
            dp[i][0] = 0
            dp[i][1] = -prices[i]
            continue
        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[n-1][0]

# Space complexity optimized version
def maxProfit_k_1(prices: List[int]) -> int:
    n = len(prices)
    # base case: dp[-1][0] = 0, dp[-1][1] = -infinity
    dp_i_0, dp_i_1 = 0, float('-inf')
    for i in range(n):
        # dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
        dp_i_0, dp_i_1 = max(dp_i_0, dp_i_1 + prices[i]), max(dp_i_1, -prices[i])
    return dp_i_0

// Original version
func maxProfit_k_1(prices []int) int {
    n := len(prices)
    dp := make([][2]int, n)
    for i := 0; i < n; i++ {
        if i - 1 == -1 {
            // base case
            dp[i][0] = 0
            dp[i][1] = -prices[i]
            continue
        }
        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[n - 1][0]
}

// Space complexity optimized version
func maxProfit_k_1(prices []int) int {
    n := len(prices)
    // base case: dp[-1][0] = 0, dp[-1][1] = -infinity
    dp_i_0, dp_i_1 := 0, -2147483648
    for i := 0; i < n; i++ {
        // dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i])
        // dp[i][1] = max(dp[i-1][1], -prices[i])
        dp_i_1 = max(dp_i_1, -prices[i])
    }
    return dp_i_0
}

var maxProfit_k_1 = function(prices) {
    var n = prices.length;
    var dp = new Array(n).fill(0).map(() => new Array(2).fill(0));
    for (var i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i-1][1], -prices[i]);
    }
    return dp[n - 1][0];
};

// version optimized for space complexity
var maxProfit_k_1 = function(prices) {
    var n = prices.length;
    // base case: dp[-1][0] = 0, dp[-1][1] = negative infinity
    var dp_i_0 = 0, dp_i_1 = -Infinity;
    for (var i = 0; i < n; i++) {
        // dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        // dp[i][1] = max(dp[i-1][1], -prices[i])
        dp_i_1 = Math.max(dp_i_1, -prices[i]);
    }
    return dp_i_0;
};

Both methods are the same, but this programming approach is much simpler. However, without the guidance of the state transition equation, it would be hard to understand. In subsequent problems, you can compare how to optimize the space of the dp array.

122. Best Time to Buy and Sell Stock II

The second problem is LeetCode's #122, "Best Time to Buy and Sell Stock II," which is the case where k is positive infinity:

Input: prices = [7,1,5,3,6,4]
Output: 7
Explanation: Buy on day 2 (price = 1) and sell on day 3 (price = 5), profit = 5-1 = 4.
Then buy on day 4 (price = 3) and sell on day 5 (price = 6), profit = 6-3 = 3.
Total profit is 4 + 3 = 7.

Input: prices = [1,2,3,4,5]
Output: 4
Explanation: Buy on day 1 (price = 1) and sell on day 5 (price = 5), profit = 5-1 = 4.
Total profit is 4.

Input: prices = [7,6,4,3,1]
Output: 0
Explanation: There is no way to make a positive profit, so we never buy the stock to achieve the maximum profit of 0.

The problem specifically emphasizes that you can sell on the same day, but I think this condition is redundant. If you buy and sell on the same day, the profit is obviously 0, which is the same as not making any transaction at all. The key feature of this problem is that it does not limit the total number of transactions k, which means k is effectively positive infinity.

If k is positive infinity, then k and k - 1 are essentially the same. You can rewrite the framework like this:

dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i])
dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i])
            = max(dp[i-1][k][1], dp[i-1][k][0] - prices[i])

We notice that the value of `k` in the array does not change, meaning we no longer need to record the state of `k`:
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])

Directly translating this into code:

// Original version
int maxProfit_k_inf(int[] prices) {
    int n = prices.length;
    int[][] dp = new int[n][2];
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i-1][1], dp[i-1][0] - prices[i]);
    }
    return dp[n - 1][0];
}

// Space complexity optimized version
int maxProfit_k_inf(int[] prices) {
    int n = prices.length;
    int dp_i_0 = 0, dp_i_1 = Integer.MIN_VALUE;
    for (int i = 0; i < n; i++) {
        int temp = dp_i_0;
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = Math.max(dp_i_1, temp - prices[i]);
    }
    return dp_i_0;
}

// Original version
int maxProfit_k_inf(vector<int>& prices) {
    int n = prices.size();
    vector<vector<int>> dp(n, vector<int>(2, 0));
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        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[n - 1][0];
}

// Space complexity optimization version
int maxProfit_k_inf(vector<int>& prices) {
    int n = prices.size();
    int dp_i_0 = 0, dp_i_1 = INT_MIN;
    for (int i = 0; i < n; i++) {
        int temp = dp_i_0;
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = max(dp_i_1, temp - prices[i]);
    }
    return dp_i_0;
}

# Original version
class Solution:
    def maxProfit_k_inf(self, prices: List[int]) -> int:
        n = len(prices)
        dp = [[0]*2 for _ in range(n)]
        for i in range(n):
            if i - 1 == -1:
                # base case
                dp[i][0] = 0
                dp[i][1] = -prices[i]
                continue
            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[n - 1][0]

# Space complexity optimized version
class Solution:
    def maxProfit_k_inf(self, prices: List[int]) -> int:
        n = len(prices)
        dp_i_0, dp_i_1 = 0, float('-inf')
        for i in range(n):
            temp = dp_i_0
            dp_i_0 = max(dp_i_0, dp_i_1 + prices[i])
            dp_i_1 = max(dp_i_1, temp - prices[i])
        return dp_i_0

// Original version
func maxProfit(prices []int) int {
    n := len(prices)
    dp := make([][2]int, n)
    for i := 0; i < n; i++ {
        if i - 1 == -1 {
            dp[i][0] = 0
            dp[i][1] = -prices[i]
            continue
        }
        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[n-1][0]
}

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

// Space complexity optimized version
func maxProfit(prices []int) int {
    n := len(prices)
    dp_i_0, dp_i_1 := 0, math.MinInt32
    for i := 0; i < n; i++ {
        temp := dp_i_0
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i])
        dp_i_1 = max(dp_i_1, temp - prices[i])
    }
    return dp_i_0
}

// Original version
var maxProfit_k_inf = function(prices) {
    var n = prices.length;
    var dp = Array.from(Array(n), () => Array(2).fill(0));
    for (var i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        dp[i][0] = Math.max(dp[i - 1][0], dp[i - 1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i - 1][1], dp[i - 1][0] - prices[i]);
    }
    return dp[n - 1][0];
}

// Space complexity optimized version
var maxProfit_k_inf = function(prices) {
    var n = prices.length;
    var dp_i_0 = 0, dp_i_1 = Number.MIN_VALUE;
    for (var i = 0; i < n; i++) {
        var temp = dp_i_0;
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = Math.max(dp_i_1, temp - prices[i]);
    }
    return dp_i_0;
}

309. Best Time to Buy and Sell Stock with Cooldown

The third problem is LeetCode #309 "Best Time to Buy and Sell Stock with Cooldown," where k is infinite, but there is a cooldown period after each transaction:

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

Input: prices = [1]
Output: 0

This problem is similar to the previous one, with the added constraint that you must wait one day after a sell before you can trade again. This feature can be integrated into the state transition equations from the previous problem:

dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
dp[i][1] = max(dp[i-1][1], dp[i-2][0] - prices[i])
Explanation: When choosing to buy on day i, the state transition should be from day i-2, not i-1.

Translated into code:

// Original version
int maxProfit_with_cool(int[] prices) {
    int n = prices.length;
    int[][] dp = new int[n][2];
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case 1
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        if (i - 2 == -1) {
            // base case 2
            dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
            // when i - 2 is less than 0, derive the
            // corresponding base case based on the state
            dp[i][1] = Math.max(dp[i-1][1], -prices[i]);
            //   dp[i][1] 
            // = max(dp[i-1][1], dp[-1][0] - prices[i])
            // = max(dp[i-1][1], 0 - prices[i])
            // = max(dp[i-1][1], -prices[i])
            continue;
        }
        dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i-1][1], dp[i-2][0] - prices[i]);
    }
    return dp[n - 1][0];
}

// Space complexity optimization version
int maxProfit_with_cool(int[] prices) {
    int n = prices.length;
    int dp_i_0 = 0, dp_i_1 = Integer.MIN_VALUE;
    // represents dp[i-2][0]
    int dp_pre_0 = 0;
    for (int i = 0; i < n; i++) {
        int temp = dp_i_0;
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = Math.max(dp_i_1, dp_pre_0 - prices[i]);
        dp_pre_0 = temp;
    }
    return dp_i_0;
}

// Original version
int maxProfit_with_cool(vector<int>& prices) {
    int n = prices.size();
    vector<vector<int>> dp(n, vector<int>(2));
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case 1
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        if (i - 2 == -1) {
            // base case 2
            dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i]);
            // when i - 2 is less than 0, derive the
            // corresponding base case according to the state
            dp[i][1] = max(dp[i-1][1], -prices[i]);
            //   dp[i][1] 
            // = max(dp[i-1][1], dp[-1][0] - prices[i])
            // = max(dp[i-1][1], 0 - prices[i])
            // = max(dp[i-1][1], -prices[i])
            continue;
        }
        dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i]);
        dp[i][1] = max(dp[i-1][1], dp[i-2][0] - prices[i]);
    }
    return dp[n - 1][0];
}

// Space complexity optimization version
int maxProfit_with_cool(vector<int>& prices) {
    int n = prices.size();
    int dp_i_0 = 0, dp_i_1 = INT_MIN;
    // represents dp[i-2][0]
    int dp_pre_0 = 0;
    for (int i = 0; i < n; i++) {
        int temp = dp_i_0;
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = max(dp_i_1, dp_pre_0 - prices[i]);
        dp_pre_0 = temp;
    }
    return dp_i_0;
}

# Original version
def maxProfit_with_cool(prices: List[int]) -> int:
    n = len(prices)
    dp = [[0] * 2 for _ in range(n)]
    for i in range(n):
        if i - 1 == -1:
            # base case 1
            dp[i][0] = 0
            dp[i][1] = -prices[i]
            continue
        if i - 2 == -1:
            # base case 2
            dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
            # When i = 1, dp[i-2] is not valid, so we can only transition from dp[i-1][1]
            dp[i][1] = max(dp[i-1][1], -prices[i])
            # Explanation as above
            continue
        dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
        dp[i][1] = max(dp[i-1][1], dp[i-2][0] - prices[i])
    return dp[n-1][0]

# Space complexity optimized version
def maxProfit_with_cool(prices: List[int]) -> int:
    n = len(prices)
    dp_i_0, dp_i_1, dp_pre_0 = 0, float('-inf'), 0
    for i in range(n):
        temp = dp_i_0
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i])
        dp_i_1 = max(dp_i_1, dp_pre_0 - prices[i])
        dp_pre_0 = temp
    return dp_i_0

// Original version
func maxProfit_with_cool(prices []int) int {
    n := len(prices)
    dp := make([][2]int, n)
    for i := 0; i < n; i++ {
        if i - 1 == -1 {
            // base case 1
            dp[i][0] = 0
            dp[i][1] = -prices[i]
            continue
        }
        if i - 2 == -1 {
            // base case 2
            dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
            // when i - 2 is less than 0, derive the
            // corresponding base case based on the state
            dp[i][1] = max(dp[i-1][1], -prices[i])
            //   dp[i][1] 
            // = max(dp[i-1][1], dp[-1][0] - prices[i])
            // = max(dp[i-1][1], 0 - prices[i])
            // = max(dp[i-1][1], -prices[i])
            continue
        }
        dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
        dp[i][1] = max(dp[i-1][1], dp[i-2][0] - prices[i])
    }
    return dp[n - 1][0]
}

// Space complexity optimization version
func maxProfit_with_cool(prices []int) int {
    n := len(prices)
    dp_i_0, dp_i_1 := 0, math.MinInt32
    // represents dp[i-2][0]
    dp_pre_0 := 0
    for i := 0; i < n; i++ {
        temp := dp_i_0
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i])
        dp_i_1 = max(dp_i_1, dp_pre_0 - prices[i])
        dp_pre_0 = temp
    }
    return dp_i_0
}

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

var maxProfit_with_cool = function(prices) {
    let n = prices.length;
    let dp = new Array(n).fill(null).map(() => new Array(2).fill(0));
    for (let i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case 1
            dp[i][0] = 0;
            dp[i][1] = -prices[i];
            continue;
        }
        if (i - 2 == -1) {
            // base case 2
            dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
            // when i - 2 is less than 0, derive the
            // corresponding base case based on the state
            dp[i][1] = Math.max(dp[i-1][1], -prices[i]);
            //   dp[i][1] 
            // = max(dp[i-1][1], dp[-1][0] - prices[i])
            // = max(dp[i-1][1], 0 - prices[i])
            // = max(dp[i-1][1], -prices[i])
            continue;
        }
        dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i-1][1], dp[i-2][0] - prices[i]);
    }
    return dp[n - 1][0];
};

var maxProfit_with_cool = function(prices) {
    let n = prices.length;
    let dp_i_0 = 0;
    let dp_i_1 = Number.MIN_SAFE_INTEGER;
    // represents dp[i-2][0]
    let dp_pre_0 = 0;
    for (let i = 0; i < n; i++) {
        let temp = dp_i_0;
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = Math.max(dp_i_1, dp_pre_0 - prices[i]);
        dp_pre_0 = temp;
    }
    return dp_i_0;
};

714. 买卖股票的最佳时机含手续费

Fourth Question: Look at LeetCode Problem #714 "Best Time to Buy and Sell Stock with Transaction Fee," which involves an infinite number of transactions (k is positive infinity) and includes transaction fees:

Input: prices = [1,3,2,8,4,9], fee = 2
Output: 8
Explanation: The maximum profit can be achieved by:
- Buying at prices[0] = 1
- Selling at prices[3] = 8
- Buying at prices[4] = 4
- Selling at prices[5] = 9
The total profit is ((8 - 1) - 2) + ((9 - 4) - 2) = 8.

Input: prices = [1,3,7,5,10,3], fee = 3
Output: 6

For each transaction, a fee must be paid. Simply subtract the fee from the profit to modify the equation:

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)
Explanation: This is equivalent to increasing the purchase price of the stock.
Subtracting in the first equation also works, as it is equivalent to decreasing the selling price of the stock.

Translate this directly into code, noting that the base case must also be adjusted to correspond with the changes in the state transition equation:

// Original version
int maxProfit_with_fee(int[] prices, int fee) {
    int n = prices.length;
    int[][] dp = new int[n][2];
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i] - fee;
            //   dp[i][1]
            // = max(dp[i - 1][1], dp[i - 1][0] - prices[i] - fee)
            // = max(dp[-1][1], dp[-1][0] - prices[i] - fee)
            // = max(-inf, 0 - prices[i] - fee)
            // = -prices[i] - fee
            continue;
        }
        dp[i][0] = Math.max(dp[i - 1][0], dp[i - 1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i - 1][1], dp[i - 1][0] - prices[i] - fee);
    }
    return dp[n - 1][0];
}

// Space complexity optimization version
int maxProfit_with_fee(int[] prices, int fee) {
    int n = prices.length;
    int dp_i_0 = 0, dp_i_1 = Integer.MIN_VALUE;
    for (int i = 0; i < n; i++) {
        int temp = dp_i_0;
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = Math.max(dp_i_1, temp - prices[i] - fee);
    }
    return dp_i_0;
}

// Original version
int maxProfit_with_fee(vector<int>& prices, int fee) {
    int n = prices.size();
    vector<vector<int>> dp(n, vector<int>(2));
    for (int i = 0; i < n; i++) {
        if (i - 1 == -1) {
            // base case
            dp[i][0] = 0;
            dp[i][1] = -prices[i] - fee;
            //   dp[i][1]
            // = max(dp[i - 1][1], dp[i - 1][0] - prices[i] - fee)
            // = max(dp[-1][1], dp[-1][0] - prices[i] - fee)
            // = max(-inf, 0 - prices[i] - fee)
            // = -prices[i] - fee
            continue;
        }
        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];
}

// Space complexity optimization version
int maxProfit_with_fee(vector<int>& prices, int fee) {
    int n = prices.size();
    int dp_i_0 = 0, dp_i_1 = INT_MIN;
    for (int i = 0; i < n; i++) {
        int temp = dp_i_0;
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = max(dp_i_1, temp - prices[i] - fee);
    }
    return dp_i_0;
}

# Original version
def maxProfit_with_fee(prices: List[int], fee: int) -> int:
    n = len(prices)
    dp = [[0] * 2 for _ in range(n)]
    for i in range(n):
        if i - 1 == -1:
            # base case
            dp[i][0] = 0
            dp[i][1] = -prices[i] - fee
            #   dp[i][1]
            # = max(dp[i - 1][1], dp[i - 1][0] - prices[i] - fee)
            # = max(dp[-1][1], dp[-1][0] - prices[i] - fee)
            # = max(-inf, 0 - prices[i] - fee)
            # = -prices[i] - fee
            continue
        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]

# Space complexity optimized version
def maxProfit_with_fee(prices: List[int], fee: int) -> int:
    n = len(prices)
    dp_i_0, dp_i_1 = 0, float('-inf')
    for i in range(n):
        temp = dp_i_0
        dp_i_0 = max(dp_i_0, dp_i_1 + prices[i])
        dp_i_1 = max(dp_i_1, temp - prices[i] - fee)
    return dp_i_0

// Original version
func maxProfit_with_fee(prices []int, fee int) int {
    n := len(prices)
    dp := make([][2]int, n)
    for i := 0; i < n; i++ {
        if i - 1 == -1 {
            // base case
            dp[i][0] = 0
            dp[i][1] = -prices[i] - fee
            //   dp[i][1]
            // = max(dp[i - 1][1], dp[i - 1][0] - prices[i] - fee)
            // = max(dp[-1][1], dp[-1][0] - prices[i] - fee)
            // = max(-inf, 0 - prices[i] - fee)
            // = -prices[i] - fee
            continue
        }
        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]
}

// Space complexity optimized version
func maxProfit_with_fee(prices []int, fee int) int {
    n := len(prices)
    dp_i_0, dp_i_1 := 0, -1<<31
    for i := 0; i < n; i++ {
        temp := dp_i_0
        dp_i_0 = Max(dp_i_0, dp_i_1+prices[i])
        dp_i_1 = Max(dp_i_1, temp-prices[i]-fee)
    }
    return dp_i_0
}

func Max(x, y int) int {
    if x > y {
        return x
    }
    return y
}

var maxProfit_with_fee = function(prices, fee) {
    var n = prices.length;
    var dp = new Array(n);
    for (var i = 0; i < n; i++) {
        dp[i] = new Array(2);
        if (i - 1 == -1) {
            // initial state
            dp[i][0] = 0;
            dp[i][1] = -prices[i] - fee;
            //   dp[i][1]
            // = max(dp[i - 1][1], dp[i - 1][0] - prices[i] - fee)
            // = max(dp[-1][1], dp[-1][0] - prices[i] - fee)
            // = max(-inf, 0 - prices[i] - fee)
            // = -prices[i] - fee
            continue;
        }
        dp[i][0] = Math.max(dp[i - 1][0], dp[i - 1][1] + prices[i]);
        dp[i][1] = Math.max(dp[i - 1][1], dp[i - 1][0] - prices[i] - fee);
    }
    return dp[n - 1][0];
};

var maxProfit_with_fee_2 = function(prices, fee) {
    var n = prices.length;
    var dp_i_0 = 0, dp_i_1 = Number.MIN_VALUE;
    for (var i = 0; i < n; i++) {
        var temp = dp_i_0;
        dp_i_0 = Math.max(dp_i_0, dp_i_1 + prices[i]);
        dp_i_1 = Math.max(dp_i_1, temp - prices[i] - fee);
    }
    return dp_i_0;
};

123. Best Time to Buy and Sell Stock III

Fifth question, refer to LeetCode problem 123 "Best Time to Buy and Sell Stock III," which deals with the case when k = 2:

Input: prices = [3,3,5,0,0,3,1,4]
Output: 6
Explanation: Buy on day 4 (price = 0) and sell on day 6 (price = 3), profit = 3-0 = 3.
Then buy on day 7 (price = 1) and sell on day 8 (price = 4), profit = 4-1 = 3.

Input: prices = [1,2,3,4,5]
Output: 4
Explanation: Buy on day 1 (price = 1) and sell on day 5 (price = 5), profit = 5-1 = 4.
Note that you cannot buy on day 1, buy on day 2 and sell them later, as you are engaging multiple transactions at the same time. You must sell before buying again.

Input: prices = [7,6,4,3,1]
Output: 0
Explanation: In this case, no transaction is done, i.e. max profit = 0.

When k = 2, the scenario is slightly different from previous problems. In those cases, the value of k didn't play a significant role: either k was infinite, making the state transition independent of k, or k = 1, which was close to the base case of k = 0 and didn't have much impact.

In this problem, where k = 2, and in the upcoming discussion where k is any positive integer, the handling of k becomes prominent. Let's write the code and analyze it as we go.

The original state transition equation, with no simplifications possible
dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i])
dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i])

Following the previous code, we might instinctively write the code like this (incorrectly):

int k = 2;
int[][][] dp = new int[n][k + 1][2];
for (int i = 0; i < n; i++) {
    if (i - 1 == -1) {
        // handle the base case
        dp[i][k][0] = 0;
        dp[i][k][1] = -prices[i];
        continue;
    }
    dp[i][k][0] = Math.max(dp[i-1][k][0], dp[i-1][k][1] + prices[i]);
    dp[i][k][1] = Math.max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i]);
}
return dp[n - 1][k][0];

int k = 2;
vector<vector<vector<int>>> dp(n, vector<vector<int>>(k + 1, vector<int>(2)));
for (int i = 0; i < n; i++) {
    if (i - 1 == -1) {
        // handle the base case
        dp[i][k][0] = 0;
        dp[i][k][1] = -prices[i];
        continue;
    }
    dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i]);
    dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i]);
}
return dp[n - 1][k][0];

k = 2
dp = [[[0 for _ in range(2)] for _ in range(k + 1)] for _ in range(n)]
for i in range(n):
    if i - 1 == -1:
        # handle the base case
        dp[i][k][0] = 0
        dp[i][k][1] = -prices[i]
        continue
    dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i])
    dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i])
return dp[n - 1][k][0]

k := 2
dp := make([][][]int, n)
for i := range dp {
    // build a 3-dimensional DP array
    dp[i] = make([][]int, k+1)
    for j := range dp[i] {
        // initialize the DP array
        dp[i][j] = make([]int, 2)
    }
}

for i := 0; i < n; i++ {
    if i-1 == -1 {
        // handle the base case
        dp[i][k][0] = 0
        dp[i][k][1] = -prices[i]
        continue
    }
    dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1]+prices[i])
    dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0]-prices[i])
}

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

var k = 2;
var dp = new Array(n);
for (var i = 0; i < n; i++) {
    dp[i] = new Array(k + 1);
    for (var j = 0; j < k + 1; j++) {
        dp[i][j] = new Array(2);
    }
    if (i - 1 == -1) {
        // handle the base case
        dp[i][k][0] = 0;
        dp[i][k][1] = -prices[i];
        continue;
    }
    dp[i][k][0] = Math.max(dp[i-1][k][0], dp[i-1][k][1] + prices[i]);
    dp[i][k][1] = Math.max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i]);
}
return dp[n - 1][k][0];

Why is it wrong? I followed the state transition equation, didn't I?

Do you remember the "Exhaustive Framework" we summarized earlier? It means we must enumerate all possible states. In fact, our previous solutions were all about enumerating all states, but in those problems, k was simplified.

For example, the code framework for the first problem when k = 1:

int n = prices.length;
int[][] dp = new int[n][2];
for (int i = 0; i < n; i++) {
    dp[i][0] = Math.max(dp[i-1][0], dp[i-1][1] + prices[i]);
    dp[i][1] = Math.max(dp[i-1][1], -prices[i]);
}
return dp[n - 1][0];

int n = prices.size();
vector<vector<int>> dp(n, vector<int>(2));
dp[0][0] = 0;
dp[0][1] = -prices[0];

for (int 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], -prices[i]);
}

return dp[n-1][0];

# LeeCode Solution
n = len(prices)
dp = [[0]*2 for _ in range(n)]
for i in range(n):
    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[n - 1][0]