Leetcode: Best Time to Buy and Sell Stock (Extension, DP)

Leetcode link to original problems:
BS1: https://oj.leetcode.com/problems/best-time-to-buy-and-sell-stock/ (O(n), O(1))
BS2: https://oj.leetcode.com/problems/best-time-to-buy-and-sell-stock-ii/ (O(n), O(1))
BS3: https://oj.leetcode.com/problems/best-time-to-buy-and-sell-stock-iii/ (O(n), O(n)) or (O(n), O(1))

Extension to buy and sell stock problem:

Q: Say you have an array of size n, for which the ith element is the price of a given stock on day i. Design an algorithm to find the maximum profit. You may complete at most k transactions.
Assume this array is named as prices[].

First step: Basic DP formula

Idea comes from lrd15's awesome post (link). Another idea that I found is this post (Chinese), which is quite hard to comprehend for me though.

We are now trying to handle a generalized problem that can be represented as finding profit[i][k], where the first index i represents day i, and second index k represent maximum number of transactions so far. As usual, for problems that aim to find maximum/minimum result, we break it to one subproblem that contains the current index, and one subproblem that does not contain the current index. This problem is no different.

Our problem:

For profit[i][k], we seek:

1. Subproblem 1: No transaction ends at day i. =>           profit[i-1][k] = S1
2. Subproblem 2: The last transaction ends at day i => S2

What we need to do is to find out which one is bigger? S1 or S2. Then profit[i][j] = S1>S2?S1:S2.
This S2 can be further decomposed to:

Suppose the last transaction starts at day j, then
S2 = maximum { profit[j][k] + (prices[i] - prices[j])}, where j is [0, i).
Now we can write the basic DP formula as:
profit[i][k] = max{ profit[i-1][k], maximum{profit[j][k] + (prices[i] - prices[j]) | j in [0,i)} }

A bottom-up implementation using (O(kn^2), O(kn)) time-space complexity.

int maxProfit(vector<int> &prices, int times = 2) {
 // corner cases
 if (prices.size() <= 1) return 0;

 // define profit[i][k] at i day, at most k transactions
 // 
 // the recursive equation is:
 // profit[i][k] = max(profit[i-1][k], ->last trans@today
 //       max(prices[i]-prices[j]+profit[j][k-1]) ->last trans@day j)
 //       here j is at range [0, i-1]
 // can also be rewritten as 
 // profit[i][k] = max(profit[i-1][k], 
 //                    prices[i]+max(profit[j][k-1]-prices[j]))
 int numOfDays = prices.size();
 int numOfTrans = times;

 int** profit = new int*[numOfDays];
 for (int i = 0; i < numOfDays; i++)
  profit[i] = new int[numOfTrans+1]();

 for (int i = 1; i < numOfDays; i++) {
  for (int k = 1; k <= numOfTrans; k++) {
   // find max(profit[j][k-1] - prices[j]) for j [0,i)
   int S = INT_MIN;
   for (int j = 0; j < i; j++) {
    int temp = profit[j][k - 1] - prices[j];
    if (temp>S) S = temp;
   }
   int term1 = profit[i - 1][k];
   int term2 = prices[i] + S;
   profit[i][k] = term1 > term2 ? term1 : term2;
  }
 }
 return profit[numOfDays - 1][numOfTrans];
}

Playable version here. Leetcode TLE since this is basically n^2 complexity.

Second step: Optimization

Now this is the tricky part since it's all about math. Let's take a look at the DP formula again:
profit[i][k] = max{ profit[i-1][k], maximum{profit[j][k] + (prices[i] - prices[j]) | j in [0,i)} }

When we calculate maximum{profit[j][k] + (prices[i] - prices[j]) | j in [0,i)} part, we notice prices[i] term can be extracted, changing the original DP formula to:
profit[i][k] = max{ profit[i-1][k], prices[i] + max{profit[j][k-1] - prices[j]} }
and the term max{profit[j][k-1] - prices[j]} can be further DP since this is a running maximum value problem.

If we write:
S[m][k] = max{ profit[j][k-1] - prices[j] | j in [0, m)}
then based on running maximum:
S[m][k] = max{ profit[m-1][k-1] - prices[m-1], S[m-1][k]}
Now we can decompose the modified DP formula to a DP equation set:
S[i][k] = max{ profit[i-1][k-1] - prices[i-1], S[i-1][k] }
profit[i][k] = max{ profit[i-1][k], prices[i] + S[i][k] }

This is the end of time optimization.

Space optimization is easy to see. profit[i][k] only depends on the last row. So does S[i][k]. Thus using bottom-up space optimization rule, we come up with the following code: (O(kn), O(k))

int maxProfit(vector<int> &prices, int times=2) {
    // corner cases
    if (prices.size() <= 1) return 0;

    // define profit[i][k] at i day, at most k transactions
    // 
    // the recursive equation is:
    // profit[i][k] = max(profit[i-1][k], ->last trans@today
    //                      max(prices[i]-prices[j]+profit[j][k-1]) ->last trans@day j)
    //                      here j is at range [0, i-1]
    // can also be rewritten as 
    // profit[i][k] = max(profit[i-1][k], 
    //                    prices[i]+max(profit[j][k-1]-prices[j]))
    // if we write:
    // S[l][k] = max(profit[j][k-1]-prices[j]), (j in [0, l)), then its equation is
    // S[l][k] = max(profit[l-1][k-1]-prices[l-1], S[l-1][k]).
    // Rewrite line 18 equation as equation set:
    //     S[i][k] = max(profit[i-1][k-1]-prices[i-1], S[i-1][k])
    //     profit[i][k] = max(profit[i-1][k], prices[i] + S[i][k])
    //
    int numOfDays = prices.size();
    int totalTransTime = times;
    // only need to construct one row
    // at day 0, no matter how many trans its all 0 profit.
    int* profit_i = new int[totalTransTime+1]();
    int* S = new int[totalTransTime+1]();
    for (int i = 0; i < totalTransTime + 1; i++)
        S[i] = INT_MIN;
    for (int i = 1; i < numOfDays; i++) {
        for (int k = totalTransTime; k > 0; k--) {
            int cand = profit_i[k - 1] - prices[i - 1];
            if (cand > S[k]) {
                S[k] = cand;
            }
            else
                S[k] = S[k];
                
            int term1 = profit_i[k];
            int term2 = prices[i] + S[k];

            profit_i[k] = term1 > term2 ? term1 : term2;
        }
    }
    return profit_i[totalTransTime];
}

Playable version here. Leetcode AC, 52ms.




Update (2014-11-19)
Why this works? Especially the decomposition part where I decompose a double maximization problem into a DP equation set? (Sidenote: Using more space to store intermediate value such that time can be saved).
The problem in this post contains subproblem max{profit[j][k-1] - prices[j]} that has nothing to do with i except in the range of j.
Can I apply this strategy to other DP problems such as "container with most water problem"?

The DP equation in the container problem is:
bestContainer[i] = max{ min(height[i], height[j]) * (j - i) | j in [0, i) }
This problem, on the contrary, does not have an intermediate result that I can store. I can not separate j from i because of the term min(height[i], height[j]).

So when to apply this space vs time strategy? I guess observation is the key. Identify the characteristic of the DP formula and see what you can find.



Update (2015-10-4)
Leetcode created a OJ for this general case. Link. One more pruning is needed to pass the large data set: If transaction time is larger or equal than day size / 2, to make sure no two transactions are overlapping, this problem degenerates to "Best time to buy and sell stock II", which means unlimited sale/buy can be performed. This can be done in O(n) time.