Problem overview
The method maxProfit receives prices, where prices[i] is the stock price on day i, and an integer k. It returns the largest profit you can make with at most k transactions, where a transaction is one buy followed by one sell. You can't hold two shares at once, so a sell must come before the next buy.
For [2, 4, 1] with k = 2 the answer is 2: buy at 2, sell at 4. 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.
Approach: turn the chart into independent trades
The usual solution is dynamic programming over days and transaction counts. This lesson uses a different, greedy construction that runs in O(n log n) because of a final sort.
Any good trade buys at the bottom of a dip and sells at the top of the rise that follows. So the first job is to cut the chart into (bottom, top) pairs, one for every dip followed by a rise. In [3, 2, 6, 5, 0, 3] they are [2, 6] and [0, 3].
If k is at least the number of pairs, take them all. If k is smaller, you have to choose, and the pairs can be unequal in a way that matters. Take the pairs [1, 4] and [2, 5]. With one transaction, the best choice is not either pair alone but buying at 1 and selling at 5 for 4. With two transactions the best is both pairs, for 3 + 3 = 6.
That suggests a trick: merging two pairs (b1, t1) and (b2, t2) into (b1, t2) is equivalent to doing the two separate trades, because (t1 - b1) + (t2 - b2) equals (t2 - b1) + (t1 - b2). So after a merge, the profit t1 - b2 is kept as an extra candidate trade. Picking k candidates from the final list then gives the answer: the single merged trade when one is allowed, and the separate trades once more are allowed.
A stack does the merging. Intervals wait on the stack in a state where their bottoms increase and their tops decrease, and each new interval either settles the ones it undercuts or merges with the ones it outclimbs. The end result is a list of independent profits, and the answer is the sum of the k largest.
Step by step
Finding the bottoms and tops
- Declare
findBottomTopPricestaking an unlabelled[Int]and returning a[[Int]]of[bottom, top]pairs. - Set up
n,iand the result array.iis the day being examined and the result starts empty. - Loop while
i < n, so every day belongs to some interval. - Walk
bottomdown the dip. Starting ati, move forward while the current price is at least the next one. The bounds checkbottom + 1 < ncomes first, so the next price is never read past the end. - Walk
topup the rise. Starting atbottom, move forward while the price is at most the next one, again checking the bounds first. - Record the pair
[prices[bottom], prices[top]]and continue fromtop + 1.
Turning intervals into independent profits
- Declare
computeProfitsover the list of pairs, with an emptystackand an emptyprofitsarray. - For each interval, copy its ends.
bottomis avarbecause a merge can lower it, andtopis alet. - Settle intervals that can no longer grow. While the stack isn't empty and the new bottom is lower than the stacked one, the stacked interval's profit is final: record
top - bottomfor it and pop it. - Merge intervals the new one outclimbs. While the stack isn't empty and the new top is at least the stacked top, record
stackedTop - bottomas the extra candidate, adopt the stacked bottom as the newbottom, and pop. - Push the resulting
[bottom, top]pair. - Drain the stack after the loop, recording each leftover pair's profit.
Choosing the best k
- Call both helpers from
maxProfit. - Cap
k. Shadow the parameter withlet k = min(k, profits.count), since there are no more candidates than that. - Sort the profits in decreasing order and sum the first
k.
The solution in Swift
class Solution {
func computeProfits(_ bottomTopPrices: [[Int]]) -> [Int] {
var stack = [[Int]]()
var profits = [Int]()
for bottomTop in bottomTopPrices {
var bottom = bottomTop[0]
let top = bottomTop[1]
while !stack.isEmpty && bottom < stack.last![0] {
profits.append(stack.last![1] - stack.last![0])
stack.removeLast()
}
while !stack.isEmpty && top >= stack.last![1] {
profits.append(stack.last![1] - bottom)
bottom = stack.last![0]
stack.removeLast()
}
stack.append([bottom, top])
}
while !stack.isEmpty {
profits.append(stack.last![1] - stack.last![0])
stack.removeLast()
}
return profits
}
func findBottomTopPrices(_ prices: [Int]) -> [[Int]] {
var n = prices.count, i = 0
var bottomTopPrices = [[Int]]()
while i < n {
var bottom = i
while bottom + 1 < n && prices[bottom] >= prices[bottom + 1] {
bottom += 1
}
var top = bottom
while top + 1 < n && prices[top] <= prices[top + 1] {
top += 1
}
bottomTopPrices.append([prices[bottom], prices[top]])
i = top + 1
}
return bottomTopPrices
}
func maxProfit(_ prices: [Int], _ k: Int) -> Int {
let bottomTopPrices = findBottomTopPrices(prices)
var profits = computeProfits(bottomTopPrices)
let k = min(k, profits.count)
profits.sort(by: { $0 > $1 })
var result = 0
for i in 0..<k {
result += profits[i]
}
return result
}
}
The Swift details worth a closer look:
- Using an array as a stack.
appendpushes,lastpeeks andremoveLast()pops.lastis an optional, sostack.last![0]force-unwraps it. That is safe here only because!stack.isEmptyis tested first and&&stops at the first false operand. - Read before you pop. Each loop body uses
stack.last!beforeremoveLast(). Swapping the two lines would read a different element or crash. - Shadowing a parameter. Swift parameters are constants, so
let k = min(k, profits.count)declares a newkin the function body. On the right-hand side,kstill means the parameter. sort(by:). The closure{ $0 > $1 }returnstruewhen its first argument should come first, which gives descending order.$0and$1are the two arguments.- The range
0..<k. Withkequal to 0 the range is empty and the loop doesn't run. That is why akof zero returns 0. - Warnings.
nis declared withvarbut never changes, so the compiler suggestslet. This is only a warning. - Method calls inside the class. Swift lets
maxProfitcall its siblings by plain name, withoutself..
Complexity
Let n be the number of prices.
Time: O(n log n). Finding the pairs is one pass, O(n). Every pair is pushed once and popped once, so computeProfits is O(n) as well. The final sort of at most n profits is O(n log n). Replacing the sort with a selection of the k largest values would bring the total down to O(n).
Space: O(n). The pairs, the stack and the profits each hold at most about n entries.
Edge cases and common mistakes
- No days. An empty price array gives no pairs and no profits, so the result is 0.
- Falling prices.
[5, 4, 3]forms a single pair[3, 3]with profit 0, so the answer is 0. - Rising prices.
[1, 2, 3, 4, 5]is one long rise, so the answer is 4 for anykof at least 1. k = 0. No transactions means a profit of 0.[3, 2, 6, 5, 0, 3]withk = 0returns 0.- Huge
k. Capping withmin(k, profits.count)keeps the loop from reading past the end ofprofits. - Choosing the best pairs without merging. For
[1, 4, 2, 5]andk = 1, picking the best raw pair gives 3, but the right answer is 4. The merging step is what finds it. Withk = 2the answer is 6. - Checking the bounds after reading. Writing
prices[bottom] >= prices[bottom + 1] && bottom + 1 < nreads past the end and crashes. - Unwrapping an empty stack. Putting
!stack.isEmptyafter the use ofstack.last!crashes the program.
Testing
[2, 4, 1],k = 2gives2.[3, 2, 6, 5, 0, 3],k = 2gives7.[3, 2, 6, 5, 0, 3],k = 1gives4.[1, 4, 2, 5],k = 1gives4.[1, 4, 2, 5],k = 2gives6.[5, 4, 3],k = 2gives0.
