Another awesome problem on LeetCode deals with buying and selling stock. You are given a list of prices, are allowed to only open long positions and you must close one before opening a new one. This would be trivial (buy on low, sell on high), but you’re limited to a total of *k* buy/sell transaction pairs. You need to return the maximum profit.

At first I thought it was a dynamic programming problem. Quick peek at the tags confirmed that suspicion. I quickly realized that subproblems require to determine the maximum profit for *d* days, starting with 2 (can’t make a profit on one day because only one price per day is given). Moreover, it is quite obvious that we have to determine maximum profits for limited numbers of transactions in range *0..k.* This requires only *k* memory because you only need the profits for the previous day to compute the tomorrow profits. You need to consider, though, that you may have left your position open, so some potential profit may exist. In this case you also have to record the opening price.

In the end it looked like this:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
public int maxProfit(int k, int[] prices) { if (prices.length <= 1 || k < 1) { return 0; } if (k >= prices.length / 2) { int sum = 0; for (int i = 1; i < prices.length; ++i) { int diff = prices[i] - prices[i - 1]; if (diff > 0) { sum += diff; } } return sum; } int[] profitClosed = new int[k + 1]; int[] profitOpen = new int[k + 1]; int[] openPrice = new int[k + 1]; int lClosed = 0, lOpen = 0; // transaction count openPrice[0] = prices[0]; int max = 0; for (int i = 1; i < prices.length; ++i) { assert lClosed >= lOpen && lClosed - lOpen <= 1; int start = lClosed; if (lOpen < lClosed) { profitOpen[++lOpen] = profitClosed[lClosed]; openPrice[lOpen] = prices[i]; } else if (lClosed < k && profitOpen[lOpen] + prices[i] - openPrice[lOpen] > profitClosed[lClosed]) { profitClosed[++lClosed] = profitOpen[lOpen] + prices[i] - openPrice[lOpen]; max = Math.max(max, profitClosed[lClosed]); } for (int j = start; j >= 1; --j) { int prevProfitClosed = profitClosed[j]; if (profitOpen[j - 1] + prices[i] - openPrice[j - 1] > profitClosed[j]) { profitClosed[j] = profitOpen[j - 1] + prices[i] - openPrice[j - 1]; max = Math.max(max, profitClosed[j]); } if (prevProfitClosed > profitOpen[j] + prices[i] - openPrice[j]) { // profit with an open position plus potential profit can now be made better profitOpen[j] = prevProfitClosed; openPrice[j] = prices[i]; } } openPrice[0] = Math.min(openPrice[0], prices[i]); } return max; } |

Not very elegant, but pretty straightforward. *profitClosed[j]* is the maximum profit that may be made by today with *j* buys/sells, while leaving no open position. *profitOpen[j]* is the same thing, but with open position and *openPrice[j]* denotes the best buying price so far, so we can instantly gain *profitOpen[j] + current price − openPrice[j]* by selling today. And that’s what we do, but only if the resulting profit is more than we could gain by performing less transaction. When we do it, the number of elements in the *profitClosed* array becomes greater than the number of elements in *profitOpen,* so we immediately open a new position on the next loop.

Then we update our profits in the inner loop. We re-close an open position if we can get a better profit and we reopen if we can get a better profit *plus potential profit!* That is important because that’s what an open position is about: if we can instantly get that much profit, it would be a waste to reopen at the current price even if we can get a better profit today, it will still bite us in the future.

This solution runs from 4 to 6 ms depending on the weather on Mars and the number of holes in the cheese. For example, replacing > profitClosed[lClosed] with > 0 in the pre-loop close-another-position condition speeds it up (branch prediction?) even though we open not very profitable positions.

There are much better solutions using the same idea. This one is particularly awesome, as it does *both loops* in just four lines.

However, I didn’t particularly like the DP idea here. It felt like this problem should have a better solution, so I kept on digging. After seeing this one, I decided to do the same thing in Java:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
public int maxProfit(int k, final int[] prices) { if (prices.length <= 1 || k < 1) { return 0; } DayList days = new DayList(prices); if (k >= days.size / 2) { return days.profit; } else { days.buildTree(); } while (k < days.size / 2) { days.removeWorst(); } return days.profit; } private static class DayList { final int[] prices; final NavigableSet<Day> dayTree = new TreeSet<>(); Day head, tail; int profit = 0, size = 0; DayList(int[] prices) { this.prices = prices; if (prices[1] > prices[0]) { add(new Day(0)); } for (int i = 1; i < prices.length; ++i) { if (i == prices.length - 1) { if (prices[i] > prices[i - 1]) { add(new Day(i)); } } else if ((prices[i] > prices[i - 1]) != (prices[i + 1] > prices[i])) { add(new Day(i)); } // can't add to tree here because day.next == null int diff = prices[i] - prices[i - 1]; if (diff > 0) { profit += diff; } } } private void add(Day day) { ++size; if (tail == null) { head = tail = day; } else { day.previous = tail; tail.next = day; tail = day; } } void buildTree() { for (Day day = head; day != null; day = day.next) { day.updateProfitLoss(); dayTree.add(day); } } void removeWorst() { // remove the day that gives the least profit Day worst = dayTree.pollFirst(); profit -= worst.profitLoss; if (worst.next == null) { Day prev = worst.previous; dayTree.remove(prev); // need to reinsert because PL may change prev.next = null; prev.updateProfitLoss(); dayTree.add(prev); } else if (worst.previous == null) { Day next = worst.next; dayTree.remove(next); // ditto next.previous = null; next.updateProfitLoss(); dayTree.add(next); } else { Day prev = worst.previous; Day next = worst.next; dayTree.remove(next); // you get the picture dayTree.remove(prev); prev.next = next; next.previous = prev; next.updateProfitLoss(); prev.updateProfitLoss(); dayTree.add(prev); dayTree.add(next); } --size; } private class Day implements Comparable<Day> { int day, profitLoss; Day previous, next; Day(int day) { this.day = day; } void updateProfitLoss() { // profit loss if we remove this day profitLoss = next == null ? Math.max(0, prices[day] - prices[previous.day]) // last day : previous == null ? Math.max(0, prices[next.day] - prices[day]) // first day : (prices[day] > prices[previous.day]) != (prices[next.day] > prices[day]) // dip or peak ? Math.min(Math.abs(prices[day] - prices[previous.day]), Math.abs(prices[next.day] - prices[day])) : 0; // neither peak nor dip } @Override public int compareTo(Day that) { int pl1 = this.profitLoss; int pl2 = that.profitLoss; int dp = Integer.compare(pl1, pl2); return dp == 0 ? Integer.compare(this.day, that.day) : dp; } } } |

Not terribly concise, but pretty straightforward. The idea is that we calculate the profit that we’ll lose if we do nothing on a particular day, thus reducing the number of transactions. Then we keep on throwing away those days starting with those that give the least profit. Obviously throwing away days on monotonous intervals won’t affect the profit at all, so we don’t even add those to begin with. We are only interested in “peaks” and “valleys” (plus possibly the first and the last days).

Maybe replacing indices arrays with a linked list was a bad idea. It’s definitely worth to try arrays instead. As it is, it runs in 20 ms. Not so terribly efficient. But I don’t think switching to array will improve it that much because the main slowdown here is the tree. Even though I only add peak/valley days and do it only if transaction reduction is needed, it is still a pretty slow thing. And it still felt like not the totally right thing.

Then I saw this solution based on another one. And I must admin, it really took me a while to figure these out. Especially loop invariants. So I hereby present my own implementation of the same ideas that runs in 3 ms (if we assume quickselect is linear, which is a fair assumption since randomization provides a very high average case probability).

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
private final static Random random = new Random(); public int maxProfit(int k, int[] prices) { if (prices.length <= 1 || k < 1) { return 0; } int[] profits = new int[prices.length / 2 + 1]; // +1 because of int[] peaks = new int[prices.length / 2 + 1]; // the possible false int[] valleys = new int[prices.length / 2 + 1]; // 0-len transaction at the end int pl = 0, top = 0; for (int p = 0, v = 0; p < prices.length; ) { v = p; while (v < prices.length - 1 && prices[v + 1] <= prices[v]) { ++v; } p = v + 1; // we need p to be the peak position + 1 to terminate the outer loop while (p < prices.length && prices[p] >= prices[p - 1]) { ++p; } // We need to find v/p pairs that satisfy both v1 <= v2 and p1 <= p2. // First let's filter out those with v1 > v2. while (top > 0 && prices[valleys[top - 1]] > prices[v]) { profits[pl++] = prices[peaks[--top]] - prices[valleys[top]]; } // Now either the stack is empty or v1 <= v2. // Let's take all pairs with p1 <= p2. while (top > 0 && prices[peaks[top - 1]] <= prices[p - 1]) { profits[pl++] = prices[peaks[--top]] - prices[v]; v = valleys[top]; } // Now either the stack is empty or p1 > p2. This maintains the stack invariant // that the v/p pairs in the stack are sorted with the minimum peak on the top. // Note that they are also sorted with the maximum valley on the top, that is, // the top pair is completely included in the deeper one and so on. // Because of this stack invariant we can do the loop above. valleys[top] = v; peaks[top++] = p - 1; } while (top > 0) { profits[pl++] = prices[peaks[--top]] - prices[valleys[top]]; } if (k < pl) { quickselect(profits, 0, pl, k); } int sum = 0, end = Math.min(k, pl); for (int i = 0; i < end; ++i) { sum += profits[i]; } return sum; } private static void quickselect(int[] nums, int start, int end, int k) { int left = start, right = end - 1; while (true) { int pi = left + random.nextInt(right - left + 1); int p = nums[pi]; nums[pi] = nums[right]; nums[right] = p; int l = left; for (int r = left; r < right; ++r) { if (nums[r] >= p) { int tmp = nums[r]; nums[r] = nums[l]; nums[l++] = tmp; } } nums[right] = nums[l]; nums[l] = p; if (l - left > k - 1) { right = l - 1; } else if (l - left < k - 1) { k -= l - left + 1; left = l + 1; } else { break; } } } |

And here comes the explanation.

Consider this price chart. Let’s consider only closing prices (80 for the first day, 70 for the 2nd, 75 for the 3rd and so on). For a “bullish” day (green) the closing price is the top of the candle, for a “bearish” one it’s the bottom. If we were to aim for the maximum profit, we’d perform the following transactions:

- Days 2–6, prices 70–100, profit 30.
- Days 8–11, prices 50–80, profit 30.
- Days 13–17, prices 20–60, profit 40.
- Days 19–21, prices 40–50, profit 10.
- Days 23–27, prices 30–70, profit 40.
- Days 29–31, prices 10–100, profit 90.

The total profit is 30+30+40+10+40+90=240.

Now, we are limited to some number of transactions. If that’s only one, then the best we can get is the last transaction, that is, 90. If two, then the best is to buy on the 13th day, sell on the 27nd (profit 50), buy on the 29th and finally sell on the 31st for the total profit of 140. Note that the first transaction is not even among the list of transactions we’d perform if we weren’t limited. That is because the first three transactions (green runs on the chart) are best united into a single one. So we need to figure out which transactions to unite.

Consider the following problem. For given intervals [v1, p1], [v2, p2], v1 < p2, v2 < p2, where “v” stands for a valley and “p” for a peak, what are the possible relationships between them? There are six:

- They don’t overlap, and p1 ≤ v2. For example, days 13–15 and 15–17 with prices [20,40], [40,60].
- They don’t overlap, and v1 ≥ p2. For example, days 2–6 and 13–17 with prices [70,100], [20,60].
- They
*do*overlap, and v2 < p1. For example, days 13–17 and 23–27 with prices [20,60], [30,70]. - They overlap, and v1 < p2. For example, days 2–6 and 8–11 with prices [70,100], [50,80].
- The second interval is fully included in the first one. For example, days 13–17 and 19–21 with prices [20,60], [40, 50].
- The first interval is fully included in the second one. For example, days 23–27 and 29–31 with prices [30,70], [10,100].

I’m not totally rigorous here about the strictly-less/less-than-or-equal thing. It doesn’t matter, though, because corner cases when something is equal to something else may be handled as either—they are kind of “at the border” between the two and belong to both sets. For example, if two transactions have exactly the same price range, then it doesn’t really matter whether you think of the second one included in the first one or vice versa. Or you may even want to consider them to be overlapping.

Now we need to ask ourselves a question: if we have two transactions and are allowed to make only one, how to get the maximum profit? Let’s consider all six cases.

The first case is not really possible for adjacent transactions if we only consider “peak—valley” transactions to begin with (the end of the second transaction is not really a peak, and the beginning of the second one is not a valley). Indeed, if two intervals form a monotonous non-decreasing sequence, then why bother splitting them in two intervals at all? However, for transactions far apart, this case is still possible (although not in our example). Anyway, in this case the answer is quite obvious: just combine two transactions into a single one, buying on for v*1**,* selling for p*2.*

The second case is also obvious: just pick up the most profitable one. We can’t unite them because the starting price of the first one is greater than the selling price of the second one. If you buy on the 2nd day and sell on 17th, you’ll have 10 loss instead of any profit.

The third case is the most tricky one! Since the lowest price is *v1,* and the highest price is p*2**, *then the best is to combine them into a single transaction. This sounds like it increases the transaction count tremendously (as we have to consider all possible combinations), but in fact it doesn’t. We do a really amazing trick here: instead of considering them as two separate transactions to begin with, we instead think of them as one “long” transaction (buying at *v1**,* selling at p*2*) and one “short” transaction (selling at p*1* and buying later at v*2*), even though short transactions aren’t technically allowed by this problem! This works because we pick transactions starting with the ones giving us the most profit. In this case, the long one is guaranteed to give us more profit than the short one, so we’ll either pick the long one (without violating anything) or pick both. And since both transactions give us exactly the same profit as two long transactions we had to begin with (p1-v1+p2-v2=p2-v1+p1-v2), then it’d appear in the net result as if we performed two separate long transactions.

The fourth case is trivial: since the combined transaction is the least profitable, we just pick the most profitable one of the two. In our example they are equally profitable, though.

The fifth case is even more boring: the first transaction has both the lowest price and the highest price, so we just pick the first one.

The sixth case is the mirror image of the fifth one. Just pick the second transaction.

So we have one case when we should combine transactions unconditionally, two cases when we should choose the most profitable of the two, two cases when the most profitable one is obvious and one tricky case when we transform two long transactions into one long and one imaginary short one. Now to the algorithm.

The algorithm preserves the following invariant. When an outer loop iteration finishes, the valleys/peaks stack contains transactions that are related according to case 5 above, with the latest transaction on the top. The invariant is obviously true at the beginning since the stack is empty, and any proposition is true for the elements of an empty set (all humans that have visited other galaxies have green ears—there are none at the moment of writing, so I’m not wrong in saying that “all” (zero) of them have green ears).

The invariant is preserved by these two loops:

1 2 3 4 5 6 7 |
while (top > 0 && prices[valleys[top - 1]] > prices[v]) { profits[pl++] = prices[peaks[--top]] - prices[valleys[top]]; } while (top > 0 && prices[peaks[top - 1]] <= prices[p - 1]) { profits[pl++] = prices[peaks[--top]] - prices[v]; v = valleys[top]; } |

The first loop just pops transactions with *v1 > v2, *which corresponds to cases 2, 4 and 6. In all these cases we don’t combine transactions, so it’s fine to just pop them and consider separately. The termination condition guarantees that *v1 ≤ v2* at the end of that loop, assuming there are any elements left in the stack.

The second loop handles the tricky case 3. We pop a transaction, then generate a “short” transaction and put it into the profits array. Then we set the current valley to the one popped from the stack, so the *current* transaction now corresponds to the long one. Then we continue the loop because it may or may not be possible to combine it with the previous one and so on. Note that the termination condition guarantees that *p1 > p2* at the end of the loop, assuming there are any elements left. Assuming the invariant held true at the start of the outer loop, popping some elements could not have increased the valley at the top of the stack because it may only decrease as we go deeper.

Note that we don’t consider case 1 separately. Instead, we treat it as a special subcase of case 3 where the “short” transaction gives us negative profit. Since we’re going to pick only the highest profits anyway, this is fine. One may think that it may have negative impact on the total profit in the case where *k* is high enough to allow all transactions to complete, but in fact it may not. This is because two adjacent transactions can’t form case 1 anyway, remember? So when we have transactions like this, it means that they are separated by some transactions in the middle. And if we’re going to pick every profitable transaction, then we can’t really combine those two because we aren’t allowed to engage in multiple transactions. So this negative profit corresponds to the profit loss caused by the fact that we need to sell first in order to buy again.

Lastly, since case 5 corresponds to the items in the stack, we pick top *k* profits using randomized quickselect and sum them up.

One last note before we get to the example: the loop may generate one false peak/valley pair at the end of the input array if the prices list ends with a valley. This corresponds to a zero-length transaction giving zero profit, so instead of checking for this corner case we can allocate one more element in the profits array to hold this zero. It won’t affect the result in any way.

Now let’s see how the algorithm works on our example. Our valleys/peaks are: 70–100, 50–80, 20–60, 40–50, 30–70, 10–100.

Step 1 (70–100). The stack is empty, so the inner loops don’t execute. The stack becomes (bottom-to-top)

1 |
stack: [70,100] |

Step 2 (50–80). The first loop pops the transaction because 70 > 50 (case 4). The stack becomes empty, therefore the second loop doesn’t execute. The next transaction is pushed into the stack.

1 2 |
profits: 30 stack: [50,80] |

Step 3 (20–60). The first loop pops because 50 > 20 (case 4). The rest is just like the previous step.

1 2 |
profits: 30 30 stack: [20,60] |

Step 4 (40–50). The first loop doesn’t work because 20 < 40. The second loop doesn’t work either because 60 > 50. Case 5.

1 2 |
profits: 30 30 stack: [20,60] [40,50] |

Step 5 (30–70). The first loop pops only [40,50] because 40 > 30 (case 6), but 20 < 30. The second loop then transforms [20,60] and [30,70] into [60,30] (added to the profits) and [20,70] (becomes the current transaction). This is case 3.

1 2 |
profits: 30 30 10 30 stack: [20,70] |

Step 6 (10–100). The first loop pops because 20 > 10 (case 6). The second loop doesn’t work because the stack is now empty. The last transaction is pushed into the stack.

1 2 |
profits: 30 30 10 30 50 stack: [10,100] |

Lastly, we pop the stack and we now have

1 |
profits: 30 30 10 30 50 90 |

A total of six transactions! Now for the different values of *k* we have:

- Just pick up the last transaction.
- Pick up 50 and 90, where 50 corresponds to 20–70 produced on the 5th step (buy on the 13th, sell on the 27th).
- Pick up 50, 90 and any of the 30s, where 30s can be one of the first two transactions or the imaginary transaction that splits our 50 into 20–60 (days 13–17) and 30–70 (days 23-27), which gives 40+40 = 50+30.
- Pick up 50, 90 and any two of the 30s.
- Pick up 50, 90 and all of the 30s (meaning 50+30 now definitely means two real transactions 40+40).
- Pick all of them (the trivial unrestricted case).

Isn’t this awesome?