Problem overview
The function max_profit receives a vector prices, where prices[i] is the price on day i, and an integer k. You may complete at most k transactions, each one a buy followed by a later sell, and you must sell before buying again. The function returns the largest total profit.
For prices = [3, 2, 6, 5, 0, 3] and k = 2 the answer is 7: buy at 2 and sell at 6 for a profit of 4, then buy at 0 and sell at 3 for a profit of 3. With k = 1 the answer would be 4, the better of the two trades.
Approach: turn the chart into independent trades, then take the best k
The usual route to this problem is dynamic programming over days and transaction counts, which takes O(n·k) time. This lesson uses a different observation: every good trade buys at a valley of the chart and sells at a peak. So the chart can be reduced to a list of (valley, peak) pairs first, and the question becomes which trades to keep.
The pairs can't simply be sorted by profit, because neighbouring pairs interact. Two consecutive rises can be one longer trade or two separate ones. A stack resolves this by splitting the chart into a list of independent candidate profits, and then the answer is the sum of the k largest of them (or all of them, when there are fewer than k).
Here is the idea on prices = [1, 2, 4, 2, 5, 7, 2, 4, 9, 0]. The valley and peak pairs are (1, 4), (2, 7), (2, 9) and (0, 0). Processing them with the stack yields the candidate profits [2, 5, 8, 0]. With k = 2 the best two are 8 and 5, so the answer is 13. With k = 3 it is 15.
Step by step
Finding the valleys and peaks
- Write
find_bottom_top_prices, which takes the prices and returns a vector of(i32, i32)tuples, one pair for each dip followed by a rise. - Track the length
nand a day indexi, and keep aresultvector for the pairs. - Loop while
i < n, so every day belongs to some pair. - Slide
bottomdown the dip. Start it atiand advance it while the next price is not higher. Checkbottom + 1 < nfirst, so the comparison never reads past the end. - Slide
topup the rise, starting frombottomand advancing while the next price is not lower, again with the bounds check first. - Store the pair of prices, then continue from
top + 1.
Turning pairs into candidate profits
- Write
compute_profits, which takes the pairs and returns aVec<i32>. Keep astackof pairs and aprofitsvector. - Take each pair as
bottomandtop, withbottommutable because merging can lower it. - Settle stacked pairs that start higher. While the stack is not empty and the new
bottomis below the stacked pair's bottom, pop that pair and record its profit. The new pair starts lower, so the popped one can never be extended. - Merge stacked pairs that end lower. While the new
topis at least the stacked pair's top, the two form one bigger trade. Peek at the stacked pair(b, t), recordt - bottomas a separate candidate, setbottomtoband pop. Two trades(b, t)and(bottom, top)earn the same in total as the merged trade(b, top)plus the trade(bottom, t), so nothing is lost by keeping the second one on its own. - Push the merged pair back on the stack, since later pairs may merge with it.
- Drain the stack with
while let, recording each remaining pair's profit.
Picking the best k
- In
max_profit, compute the pairs, then the profits, which must be mutable because they get sorted. - Cap
k. Shadow it as ausizeno larger than the number of candidates. - Sort in decreasing order and sum the first
kas the final expression.
The solution in Rust
impl Solution {
fn compute_profits(bottom_top_prices: Vec<(i32, i32)>) -> Vec<i32> {
let mut stack: Vec<(i32, i32)> = Vec::new();
let mut profits: Vec<i32> = Vec::new();
for (mut bottom, top) in bottom_top_prices {
while !stack.is_empty() && bottom < stack.last().unwrap().0 {
let (b, t) = stack.pop().unwrap();
profits.push(t - b);
}
while !stack.is_empty() && top >= stack.last().unwrap().1 {
let (b, t) = *stack.last().unwrap();
profits.push(t - bottom);
bottom = b;
stack.pop();
}
stack.push((bottom, top));
}
while let Some((b, t)) = stack.pop() {
profits.push(t - b);
}
profits
}
fn find_bottom_top_prices(prices: Vec<i32>) -> Vec<(i32, i32)> {
let n = prices.len();
let mut i = 0;
let mut result: Vec<(i32, i32)> = Vec::new();
while i < n {
let mut bottom = i;
while bottom + 1 < n && prices[bottom] >= prices[bottom + 1] {
bottom += 1;
}
let mut top = bottom;
while top + 1 < n && prices[top] <= prices[top + 1] {
top += 1;
}
result.push((prices[bottom], prices[top]));
i = top + 1;
}
result
}
fn max_profit(prices: Vec<i32>, k: i32) -> i32 {
let bottom_top_prices = Self::find_bottom_top_prices(prices);
let mut profits = Self::compute_profits(bottom_top_prices);
let k = (k as usize).min(profits.len());
profits.sort_unstable_by(|a, b| b.cmp(a));
profits[..k].iter().sum()
}
}
The Rust details worth a closer look:
- Destructuring in
for.for (mut bottom, top) in bottom_top_pricestakes each tuple apart, andmutapplies only tobottom. Looping over the vector by value moves it, which is fine because it isn't needed afterwards. last()returns anOption. It gives a reference to the final element, orNonefor an empty stack. Theunwrap()calls are safe only because!stack.is_empty()is checked first, and&&stops at the first false condition..0and.1read the tuple's fields.- Copying out of a reference.
*stack.last().unwrap()dereferences the reference. A tuple of integers isCopy, sobandtare independent values and the stack is untouched untilpop(). The popped value is ignored. while letloops for as long aspop()returnsSome, so it empties the stack and stops atNone.- Index arithmetic with
usize. The scan usesbottom + 1 < nrather thanbottom < n - 1. With an emptypricesthe second form would subtract from zero and wrap in a build with optimizations, or panic in a debug build. Here the outer loop never runs for an empty vector. - Sorting.
sort_unstable_by(|a, b| b.cmp(a))compares the second element with the first, which puts the largest values first. Equal integers are interchangeable, so an unstable sort loses nothing. - Summing a slice.
profits[..k]is the firstkelements, and.iter().sum()adds them. The function'si32return type tellssumwhat to produce. - Associated functions. The helpers take no
self, so they are called asSelf::find_bottom_top_pricesandSelf::compute_profits.
Complexity
Time: O(n log n). Finding the pairs is one pass over the prices. Each pair is pushed onto the stack once and popped at most once, so compute_profits is linear. The sort of up to n candidates dominates with O(n log n).
Space: O(n). The pairs, the stack and the profits each hold at most one entry per pair.
Edge cases and common mistakes
- Prices that only fall.
[5, 4, 3, 2, 1]gives a single pair(1, 1)with profit 0, so the answer is 0 for anyk. - Empty prices. The outer loop never runs, there are no pairs and no profits, and the sum of an empty slice is 0.
klarger than the number of candidates. Without the cap,profits[..k]would panic with an out-of-range slice. For[3, 2, 6, 5, 0, 3]andk = 5the cap leaves 2 candidates, and the answer is still 7.k = 0. The sliceprofits[..0]is empty, so the answer is 0.- Taking the best k pairs without merging. Sorting the raw valley and peak pairs by profit gives wrong answers when a deeper valley or a higher peak makes neighbouring rises combine. For the 10-day example the raw pairs have profits 3, 5, 7 and 0, so the best two would add up to 12, not 13.
- Forgetting the bounds check order. Writing
prices[bottom] >= prices[bottom + 1] && bottom + 1 < nreads the next price before checking that it exists.
Testing
prices = [2, 4, 1],k = 2gives 2.prices = [3, 2, 6, 5, 0, 3],k = 2gives 7, and withk = 1gives 4.prices = [1, 2, 4, 2, 5, 7, 2, 4, 9, 0],k = 2gives 13, and withk = 3gives 15.prices = [5, 4, 3, 2, 1],k = 2gives 0.prices = [],k = 2gives 0.
