Problem overview
The function maxProfit receives a slice prices, where prices[i] is the price of a stock on day i, and an integer k. You may buy and sell at most k times, and you must sell before you buy again. The function returns the largest total profit.
For [3, 2, 6, 5, 0, 3] with k = 2 the answer is 7: buy at 2 and sell at 6 for 4, then buy at 0 and sell at 3 for 3. With [2, 4, 1] and k = 2 the answer is 2, since only one profitable trade exists.
Approach: turn the chart into candidate trades
The usual solution is a table indexed by day and by number of transactions, which takes O(n * k) time. This lesson avoids the table with a greedy idea built on a stack, and the cost is dominated by one sort.
Start from the shape of the chart. A good trade buys at the bottom of a dip and sells at the top of the following rise, so the chart breaks into (bottom, top) pairs, one per dip-then-rise.
With enough transactions you would take every pair. With fewer, you must merge neighbouring pairs, and merging is where it gets interesting. Take two pairs (b1, t1) and (b2, t2) in time order. If b2 >= b1 and t2 >= t1, you could replace them by one long trade (b1, t2). That earns t2 - b1, and the two separate trades earn (t1 - b1) + (t2 - b2). Those totals differ by exactly t1 - b2, so the pair of separate trades equals the merged trade plus a spare candidate trade worth t1 - b2. The stack keeps producing such independent profits, each a trade you may take or skip. At the end, the best k of them are the answer.
Take [1, 4, 2, 7, 3, 5]. Its pairs are (1, 4), (2, 7) and (3, 5). The first two merge into (1, 7) worth 6, leaving a spare profit of 4 - 2 = 2. The last pair stays as it is, worth 2. The candidates are 6, 2 and 2, so the answer is 6 for k = 1, 8 for k = 2 and 10 for k = 3.
Step by step
Extracting the pairs
- Write
findBottomTopPrices. It takespricesand returns a[][2]int: a slice of two-element arrays, each holding a bottom price and a top price. - Keep a day index
iand the lengthn, and a nil slicebottomTopPricesdeclared withvarto collect pairs. - Loop while
i < n. Each pass consumes one dip and the rise after it. - Find the bottom. Start at
iand move forward while the next price is not higher. Checkbottom+1 < nfirst, so the next price is never read past the end. - Find the top. Start at the bottom and move forward while the next price is not lower, with the same bounds check first.
- Append the pair and resume after it. Set
i = top + 1, then return the slice.
Turning pairs into profits
- Write
computeProfits, which takes the pairs and returns a[]int, with an emptystackof pairs and a nilprofitsslice. - Range over the pairs and copy each into
bottomandtop. - First inner loop. While the stack is non-empty and the current
bottomis lower than the stacked pair's bottom, the stacked pair can never be extended to the right by merging. Pop it and record its profitt[1]-t[0]. - Second inner loop. While the stack is non-empty and the current
topis at least the stacked pair's top, the two merge. Record the spare profitt[1]-bottom, take over the lower buy price withbottom = t[0], and pop. - Push the current pair. A later pair may still merge with it.
- After the loop, pop what is left, recording each pair's own profit, then return
profits.
Choosing the best k
- In
maxProfit, call the two helpers. - Cap
katlen(profits): you can't take more trades than there are candidates. - Sort the profits in decreasing order, add the first
kand return the total.
The solution in Go
func findBottomTopPrices(prices []int) [][2]int {
n := len(prices)
i := 0
var bottomTopPrices [][2]int
for i < n {
bottom := i
for bottom+1 < n && prices[bottom] >= prices[bottom+1] {
bottom++
}
top := bottom
for top+1 < n && prices[top] <= prices[top+1] {
top++
}
bottomTopPrices = append(bottomTopPrices, [2]int{prices[bottom], prices[top]})
i = top + 1
}
return bottomTopPrices
}
func computeProfits(bottomTopPrices [][2]int) []int {
var stack [][2]int
var profits []int
for _, bt := range bottomTopPrices {
bottom, top := bt[0], bt[1]
for len(stack) > 0 && bottom < stack[len(stack)-1][0] {
t := stack[len(stack)-1]
profits = append(profits, t[1]-t[0])
stack = stack[:len(stack)-1]
}
for len(stack) > 0 && top >= stack[len(stack)-1][1] {
t := stack[len(stack)-1]
profits = append(profits, t[1]-bottom)
bottom = t[0]
stack = stack[:len(stack)-1]
}
stack = append(stack, [2]int{bottom, top})
}
for len(stack) > 0 {
t := stack[len(stack)-1]
profits = append(profits, t[1]-t[0])
stack = stack[:len(stack)-1]
}
return profits
}
func maxProfit(prices []int, k int) int {
bottomTopPrices := findBottomTopPrices(prices)
profits := computeProfits(bottomTopPrices)
if k > len(profits) {
k = len(profits)
}
sort.Sort(sort.Reverse(sort.IntSlice(profits)))
total := 0
for i := 0; i < k; i++ {
total += profits[i]
}
return total
}
The Go details worth knowing:
[2]intis an array, not a slice. Its length is part of its type, and arrays are values, sot := stack[len(stack)-1]copies the pair. That is whytstays valid after the stack shrinks.- A nil slice works with
append.var stack [][2]intneeds nomake, andappendreturns the grown slice, so its result must be assigned back. - Push and pop are slice operations. Push is
append. Pop is reslicing withstack[:len(stack)-1], which drops the last element without copying. - Go has no
while.for i < nandfor len(stack) > 0are the condition-only form offor. - Short-circuit
&&. Inbottom+1 < n && prices[bottom] >= prices[bottom+1]the right side is not evaluated when the bounds test fails. The same holds for the emptiness test beforestack[len(stack)-1]. Reversing either order would panic with an index out of range. - Sorting descending uses
sort.Sort(sort.Reverse(sort.IntSlice(profits))).sort.IntSliceturns the slice into something thesortpackage can order,sort.Reverseflips the comparison, and the sort happens in place. The solution needsimport "sort".
Complexity
Extracting the pairs is O(n), because i only moves forward. In computeProfits every pair is pushed once and popped once, so it is O(n) in total even with loops inside the loop. Sorting the profits costs O(n log n), which dominates. Space is O(n) for the pairs, the stack and the profits. Picking the top k without a full sort can bring the time down to O(n).
Edge cases and common mistakes
- Falling prices. For
[5, 4, 3]the only pair is(3, 3), worth 0, so the result is 0. The solution never returns a negative profit. - Empty or one-day input. No trade is possible. An empty slice gives no pairs, so
profitsis empty and the loop overkdoesn't run. - Large
k. Withk = 5on the first example the cap reduces it to the number of candidates, and the result is 7, the same as withk = 2. Without the cap,profits[i]would panic. k = 0. The loop adds nothing and the result is 0.- Using
>instead of>=in the scans. Equal neighbouring prices would then end a dip or a rise early and create extra pairs. - Forgetting to assign the result of
append. The stack or the profits would silently stay empty. - Selling and buying on the same day. Consecutive pairs are separate trades, and each sell happens before the next buy.
Testing
Inputs worth trying, with the results the solution gives:
[]int{2, 4, 1}withk = 2returns 2.[]int{3, 2, 6, 5, 0, 3}withk = 2returns 7, and withk = 1returns 4.[]int{1, 4, 2, 7, 3, 5}withk = 1,2,3returns 6, 8, 10.[]int{1, 2, 3, 4, 5}withk = 2returns 4.[]int{5, 4, 3}withk = 2returns 0.