A General Solution to Stock Problems in Dynamic Programming

Language:中文/EN

A General Solution to Stock Problems in Dynamic Programming

There is a category of dynamic programming problems where you are given a sequence of stock prices and need to calculate the maximum profit that can be obtained from buying and selling stocks. These problems often have many variations, such as allowing only one transaction, multiple transactions, or including transaction fees. The maximum profit usually depends on the timing of the transactions and the maximum number of allowed transactions (each transaction consists of one buy and one sell).

We can use T[i][k] to represent the maximum profit that can be obtained by the end of day i with at most k transactions. Additionally, by the end of day i, there can be two states: having a stock T[i][k][1] or not having a stock T[i][k][0]. The initial states can be defined as follows:

T[-1][k][0] = 0, T[-1][k][1] = -Infinity
T[i][0][0] = 0, T[i][0][1] = -Infinity

Here, T[-1][k][0] = 0 and T[i][0][0] = 0 mean that the profit is 0 when there are no stocks at the initial state (i=-1) or when 0 transactions are allowed. T[-1][k][1] = T[i][0][1] = -Infinity means that having a stock in hand at the initial state or with 0 allowed transactions is impossible, so the profit is set to -Infinity.

For state transitions, we can derive the following formulas:

T[i][k][0] = max(T[i-1][k][0], T[i-1][k][1] + prices[i])
// By the end of day i, if there are no stocks, there are two possible actions:
// 1. Do not trade, i.e., the maximum profit is T[i-1][k][0]
// 2. Sell, i.e., the maximum profit on day i is the maximum profit from the previous day with a stock plus the selling price on day i, which is T[i-1][k][1] + prices[i]
T[i][k][1] = max(T[i-1][k][1], T[i-1][k-1][0] - prices[i])
// By the end of day i, if there is a stock, there are also two possible actions:
// 1. Do not trade, i.e., T[i-1][k][0]
// 2. Buy, which increases the transaction count by one, so the maximum transaction count by the end of day i-1 is k-1, i.e., the maximum profit by the end of day i is T[i-1][k-1][0] - prices[i]

Since the maximum profit is achieved when the last step is to sell the stock, the final return should be T[i][k][0].

Best Time to Buy and Sell Stock
This problem asks for the maximum profit that can be obtained with only one transaction, i.e., when k=1. The solution is as follows:

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int T_i10=0,T_i11=INT_MIN;
        for(int price:prices){
            T_i10=max(T_i10,T_i11+price);
            T_i11=max(T_i11,-price);
        }
        return T_i10;
    }
};

Best Time to Buy and Sell Stock II
This problem allows unlimited transactions, i.e., k=Infinity. In this case, k-1=k, and the state transition equations can be written as follows:

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

Since T[i][k][0] is updated in the previous step, a temporary variable can be used to store T[i][k][0]. The code is as follows:

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int T_ik0=0,T_ik1=INT_MIN;
        for(int price:prices){
            int T_ik0_old=T_ik0;
            T_ik0=max(T_ik0,T_ik1+price);
            T_ik1=max(T_ik1,T_ik0_old-price);
        }
        return T_ik0;
    }
};

Best Time to Buy and Sell Stock III
This problem seeks the maximum profit with a maximum of 2 transactions, i.e., k=2 or k=1. Similar to the k=1 case, the state transition equations can be written as follows:

T[i][2][0] = max(T[i-1][2][0], T[i-1][2][1] + prices[i])
T[i][2][1] = max(T[i-1][2][1], T[i-1][1][0] - prices[i])
T[i][1][0] = max(T[i-1][1][0], T[i-1][1][1] + prices[i])
T[i][1][1] = max(T[i-1][1][1], -prices[i])

The optimal solution to return is T[i][k][0]=T[i][2][0]. The complete solution is as follows:

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int T_i10=0,T_i11=INT_MIN;
        int T_i20=0,T_i21=INT_MIN;
        for(int price:prices){
            T_i20=max(T_i20,T_i21+price);
            T_i21=max(T_i21,T_i10-price);
            T_i10=max(T_i10,T_i11+price);
            T_i11=max(T_i11,-price);
        }
        return T_i20;
    }
};

Best Time to Buy and Sell Stock IV
In this problem, k is passed as a parameter to the solution function. It can be observed that each profitable transaction requires two days, so when k>=n/2, increasing k further does not increase profit, and it can be treated as k=Infinity, similar to the previous problem. When k < n/2, it is similar to the k=2 case, and forward updating can be done using arrays. The solution is as follows:

class Solution {
public:
    int maxProfit(int k, vector<int>& prices) {
        if(k>=prices.size()>>1){
            int T_ik0=0,T_ik1=INT_MIN;
            for(int price:prices){
                int T_ik0_old=T_ik0;
                T_ik0=max(T_ik0,T_ik1+price);
                T_ik1=max(T_ik1,T_ik0_old-price);
            }
            return T_ik0;
        }
        vector<int> T_ik0(k+1,0);
        vector<int> T_ik1(k+1,INT_MIN);
        for(int price:prices){
            for(int j=k;j>0;j--){
                T_ik0[j]=max(T_ik0[j],T_ik1[j]+price);
                T_ik1[j]=max(T_ik1[j],T_ik0[j-1]-price);
            }
        }
        return T_ik0[k];
    }
};

Best Time to Buy and Sell Stock with Cooldown
This problem does not limit the value of k, i.e., it is Infinity, but there is another restriction: you cannot buy on the day immediately after selling a stock, implementing a cooldown. In this state transition equation, the update for T[i][k][1] differs. On day i, if resting, it is the same, but if buying, it updates to T[i-2][k-1][0]=T[i-2][k][0]. The state transition equations are as follows:

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

The code for this problem is:

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int T_ik0_pre=0,T_ik0=0,T_ik1=INT_MIN;
        for(int price:prices){
            int T_ik0_old=T_ik0;
            T_ik0=max(T_ik0,T_ik1+price);
            T_ik1=max(T_ik1,T_ik0_pre-price);
            T_ik0_pre=T_ik0_old;
        }
        return T_ik0;
    }
};

Best Time to Buy and Sell Stock with Transaction Fee
This problem differs from the previous ones by adding a transaction fee for each trade, while also allowing unlimited transactions. The transaction fee can be subtracted during each trade (either selling or buying) before updating the respective maximum profit T[i][k][0] or T[i][k][1]. The state transition equations are as follows:

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

or

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

The algorithm code is as follows:

#include<climits>
class Solution {
public:
    int maxProfit(vector<int>& prices, int fee) {
        long T_ik0=0,T_ik1=LONG_MIN;
        for(int price:prices){
            long T_ik0_old=T_ik0;
            T_ik0=max(T_ik0,T_ik1+price);
            T_ik1=max(T_ik1,T_ik0_old-price-fee);
        }
        return (int)T_ik0;
    }
};

Reference article: Most consistent ways of dealing with the series of stock problems

#Algorithm