Best Time to Buy and Sell Stock IV

On a whim, I decided to solve this leetcode question.

There is a really nice solution that solves this problem in \(O(N)\) worst case. I didn't think of it, I did an \(O(Nk)\) solution. It took me a while to get so I wanted to record my thoughts.

The Question

Given an array prices of length \(N\) representing the prices of the stock on different days. Find the maximum profit by buying and selling stocks on different days. You may complete at most \(k\) transactions (up to \(k\) buys and \(k\) sells). Your position of the stock may not exceed 1 (i.e. you must sell the stock before you buy again) or go below \(0\).

Solution

Lemma 1 Consider \(prices[i]\) such that \(prices[i-1] < prices[i]\), in any optimal solution we do not buy on day \(i\) (we do nothing or sell). Symmetrically, if \(price[i] < price[i+1]\) then in any optimal solution we do not sell on day \(i\).

Proof

Lemma 2 If \(prices[i] > prices[i+1]\), then in any optimal solution we do not buy on day \(i\). Symmetrically if \(prices[i-1] > prices[i]\), then in any optimal solution we do not sell on day \(i\).

Proof

Key Observation Number 1 (Buy Low, Sell High) In any optimal solution, we buy only on days \(i\) when \(prices[i-1] \geq prices[i] \leq prices[i+1]\) (taking \(prices[0]\) to be \(\infty\) and \(price[prices.length]\) to be \(\infty\)).
Symmetrically, in any optimal solution, we sell only on days \(i\) when \(prices[i-1] \leq prices[i] \geq prices[i+1]\) (taking \(prices[0]\) to be \(-\infty\) and \(price[prices.length]\) to be \(-\infty\)).

Proof

Lemma 3 There exists an optimal solution where if \(prices[i]=prices[i+1]\) we do nothing on day \(i\). Informally, we can ignore consecutive duplicate days.

Proof

Refactoring the Input First ignore the consecutive days with duplicate prices. Construct a list of days we can buy and sell on based on key observation 1 (these lists are disjoint). We can't sell before we buy or have a buy the stock at the end without selling so we can ignore these days. The days we can buy and sell on alternate, follows from the 3rd observation. Now pair up potential buy days with the potenial sell day following directly after them.
Denote the set of pairs of buy and sell prices \(\{(b_1,s_1), (b_2,s_2), \dots, (b_m,s_m)\}\). Our question is find \(m(m \leq k)\) pairs of the form \(\{(b_{i_1},s_{j_1}),(b_{i_2},s_{j_2}), \dots, (b_{i_m},s_{j_m})\}\) where \( i_x \leq j_x \) and \(j_x < i_{x+1}\) for all valid x.
Key Observation Number 2 Let's say \(b_i \geq b_{i+1} \) and before \(b_i\) are intervals that completely cover it then in the optimal solution we either

buy at b_i and sell at \(s_i\). Furthermore, removing \((b_i,s_i)\) from our solution we will get an optimal solution with k-1 transactions for the remaining input.
do not buy at b_i and do not sell at \(s_i\)

Proof

Key Observation Number 3 Let's say \( b_{i} < b_{i+1} \text{ and } s_i \leq s_{i+1}\) and before \((b_i,s_i)\) are intervals that completely covers it. The optimal solution to input \(\{(b_1,s_1), (b_2,s_2), \dots,(b_i,s_{i}), (b_{i+1},s_{i+1}) , \dots, (b_m,s_m)\}\) is equal to the optimal solution to input \(\{(b_1,s_1), (b_2,s_2), \dots, (b_{i+1},s_{i}), (b_i,s_{i+1}), \dots, (b_m,s_m)\}\)

Proof

Consider an optimal solution to \(\{(b_1,s_1), (b_2,s_2), \dots,(b_i,s_{i}), (b_{i+1},s_{i+1}) , \dots, (b_m,s_m)\}\). We will either have 0 position of the stock entering \((b_i,s_{i})\) or +1 position of the stock. If we enter \((b_i,s_i)\) with a +1 position then we will never sell at \(s_i\) as we could of sold earlier for at least as good a price, but acting optimally we can either

sell at \(s_{i+1}\)
do nothing on days with price \(b_i\), \(s_i\), \(b_{i+1}\), \(s_{i+1}\)

If we enter \((b_i,s_i)\) with a 0 position then acting optimally we will either

buy at \(b_i\), sell at \(s_{i}\), buy at \(b_{i+1}\)
buy at \(b_i\), sell at \(s_i\), buy at \(b_{i+1}\), and sell at \(s_{i+1}\)
buy at \(b_i\), sell at \(s_{i+1}\)
buy at \(b_i\)
do nothing on days with price b_i, \(s_i\), b_{i+1}, \(s_{i+1}\)

while doing nothing on days not specified. Can show all other actions can be changed into one of these actions.

Now consider if the input was instead \(\{(b_1,s_1), (b_2,s_2), \dots, (b_{i+1},s_{i}), (b_i,s_{i+1}), \dots, (b_m,s_m)\}\) If we enter \((b_i,s_i)\) with a +1 position then we will never sell at \(s_i\) as we could of sold earlier for as good a price. Acting optimally we can either

sell at \(s_{i+1}\)
do nothing on days with price \(b_i\), \(s_i\), b_{i+1}, \(s_{i+1}\)

If we enter \((b_i,s_i)\) with a 0 position. Acting optimally we can either

buy at \(b_{i+1}\), sell at \(s_{i}\), buy at \(b_{i}\)
buy at \(b_i\), sell at \(s_i\), buy at \(b_{i+1}\), and sell at \(s_{i+1}\)
buy at \(b_i\), sell at \(s_{i+1}\)
buy at \(b_i\)
do nothing on days with price \(b_i\), \(s_i\), b_{i+1}, \(s_{i+1}\)

while doing nothing on days not specified. Can show all other actions can be changed into one of these actions. Because our optimal set of actions result in the same outcome solving either problem results in the same solution.

Solving the problem Key Observation 2+3 suggest a interesting approach. We can iterate through \( (b_i, s_i) \).
If it has buy price is less than or equal to the previous buy price (\(b_i \leq b_{i-1}\)), then by observation two we can remove it and solve the problem on the remaining buy sell pairs. We can determine the optimal solution with \(k\) transactions and \(k-1\) transactions and if that difference is less than \(s_i-b_i\) then we should include it along with the other optimal buy sell pairs for \(k-1\) transactions.
If it has a buy price greater than the last one and sell price greater than or equal to the last one \( ( b_{i-1} < b_i \text{ and } s_{i-1}\leq s_i) \) then we can change our problem by observation three. In other words, change \(\{\dots, (b_{i-1},s_{i-1}), (b_i,s_i), \dots\}\) to \(\{\dots, (b_{i},s_{i-1}), (b_{i-1},s_i), \dots\}\) and then apply observation 2 to reduce the question again.
In this manner when iterating on \( (b_i, s_i) \), we maintain that we are always solving a problem where all the intervals before completely contain the next. When we are finished we are left with intervals each one which contains the next. To find the best solution in this particular case, with at most \(k\) transactions it suffices to just take the first \(k\) buys and first \(k\) sells (or as many as possible if less than \(k\) pairs left). Also add back the bonus generated by the intervals removed by the repeated use of observation 2. You can find code implementation here. (Instead of using a max heap use a \(O(n)\) k-selection algorithm like median of medians).

Final Thoughts The also allows you to find maximum total sum of k disjoint subarrays in an array \(a\) in \(O(n)\) where \(n\) is the length of the array. Just take the prefix array and apply this problem to it (buying is the start of a subarray and selling is the end of a subarray). In the case \(k=1\) this is the well known maximum sum subarray problem.
Define function \(ans(x)\) as the answer to this problem with at most \(x\) transactions (\(k=x\)). This reasoning and implementation proves \(ans(x)\) a concave function. You could use alien's trick.