LESSON 332 OF 348 · GREEDY

Best Time to Buy Stock in Kotlin: At Most k Transactions

Say you have an array for which the i-th element is the price of a given stock on day i.

Design an algorithm to find the maximum profit. You may complete at most k transactions.

Note:
You may not engage in multiple transactions at the same time (ie, you must sell the stock before you buy again).

Example 1:

Input: [2,4,1], k = 2
Output: 2
Explanation: Buy on day 1 (price = 2) and sell on day 2 (price = 4), profit = 4-2 = 2.

Example 2:

Input: [3,2,6,5,0,3], k = 2
Output: 7
Explanation: Buy on day 2 (price = 2) and sell on day 3 (price = 6), profit = 6-2 = 4.
             Then buy on day 5 (price = 0) and sell on day 6 (price = 3), profit = 3-0 = 3.

STEP 1 OF 29

Guided

Ctrl + Enter to check

KOTLIN · GREEDY

Best Time To Buy Stock in Kotlin, explained

Problem overview

The function maxProfit receives prices, an IntArray where prices[i] is the price of a stock on day i, and an integer k. You may complete at most k transactions, a transaction being one buy followed by one sell, 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 a profit of 4, then buy at 0 and sell at 3 for a profit of 3. With k = 1 the same prices give 4, because only the best single trade counts.

Approach: turn the chart into candidate trades

A dynamic program over days and transaction counts works, but this lesson uses a greedy decomposition that finishes after a sort.

Every worthwhile trade buys at a local minimum of the chart and sells at a local maximum. So the first job is to reduce the chart to its valleys and peaks: a list of (bottom, top) pairs, one for each dip followed by a rise. For [3, 2, 6, 5, 0, 3] that list is (2, 6) and (0, 3).

If there were as many allowed trades as pairs, you would take every pair. When k is smaller, some pairs have to be merged, and merging is where the difficulty lies. Take [1, 5, 2, 6], whose pairs are (1, 5) and (2, 6). With two trades you earn 4 + 4 = 8. With one trade you buy at 1 and sell at 6 for 5. Notice that 5 + 3 = 8: the one big trade (1, 6) plus the profit you give up by merging, 5 - 2 = 3, adds up to the two separate trades.

That gives the plan. Walk through the pairs with a stack and produce a list of independent candidate profits: the profit of every merged trade, plus the profit given up each time two pairs are merged. The best answer for k trades is the sum of the k largest candidates. The stack handles the other shape of overlap too, where a later pair starts lower than an earlier one: then the earlier pair can never be extended and its profit is final.

Sorting the candidates costs O(n log n). A selection technique could bring that to O(n), but sorting is simpler and is what the solution uses.

Step by step

Finding the bottoms and tops

Write findBottomTopPrices. It takes the prices and returns a List<IntArray>, where each two-element array holds a bottom price and a top price.
Set up the scan. n is the number of days and never changes, so it is a val. i is the current day and is a var. Create an empty mutable list for the pairs.
Loop while i < n, extracting one dip and its rise per iteration.
Find the bottom. Start at i and move forward while the price doesn't increase. The loop checks bottom + 1 < n first so that the next price is never read past the end of the array.
Find the top. Start from the bottom and move forward while the price doesn't decrease, with the same bounds check.
Store the prices at those two days as a pair.
Resume after the top by setting i = top + 1, and return the list when the loop ends.

Computing the candidate profits

Start maxProfit. Call findBottomTopPrices and keep the pairs.
Write computeProfits. It takes the list of pairs and returns a MutableList<Int> of profits. Create an empty ArrayDeque<IntArray> to act as a stack, whose last element is its top, and an empty list for the profits.
Loop over the pairs. Copy the current bottom into a var, since merging can lower it, and the top into a val.
Settle pairs that can't grow. While the stack is non-empty and the current bottom is lower than the stacked pair's bottom, add that pair's profit and pop it. The current pair starts deeper, so any later merge would use it instead.
Merge pairs that the current one outgrows. While the stack is non-empty and the current top is at least the stacked top, record the profit given up (the stacked top minus the current bottom), take the stacked pair's lower bottom, and pop it.
Push the current pair with addLast.
Empty the stack at the end. Every pair left is final, so add each one's profit.

Choosing the best k

Get the profits from computeProfits.
Cap k. A parameter is read-only, so a new val k holds the smaller of k and the number of candidates.
Sort the profits in descending order, add the first k values and return the sum.

The solution in Kotlin

class Solution {

    fun computeProfits(bottomTopPrices: List<IntArray>): MutableList<Int> {
        val stack = ArrayDeque<IntArray>()
        val profits = mutableListOf<Int>()
        for (bottomTop in bottomTopPrices) {
            var bottom = bottomTop[0]
            val top = bottomTop[1]
            while (stack.isNotEmpty() && bottom < stack.last()[0]) {
                profits.add(stack.last()[1] - stack.last()[0])
                stack.removeLast()
            }
            while (stack.isNotEmpty() && top >= stack.last()[1]) {
                profits.add(stack.last()[1] - bottom)
                bottom = stack.last()[0]
                stack.removeLast()
            }
            stack.addLast(intArrayOf(bottom, top))
        }
        while (stack.isNotEmpty()) {
            profits.add(stack.last()[1] - stack.last()[0])
            stack.removeLast()
        }
        return profits
    }

    fun findBottomTopPrices(prices: IntArray): List<IntArray> {
        val n = prices.size; var i = 0
        val bottomTopPrices = mutableListOf<IntArray>()
        while (i < n) {
            var bottom = i
            while (bottom + 1 < n && prices[bottom] >= prices[bottom + 1]) {
                ++bottom
            }
            var top = bottom
            while (top + 1 < n && prices[top] <= prices[top + 1]) {
                ++top
            }
            bottomTopPrices.add(intArrayOf(prices[bottom], prices[top]))
            i = top + 1
        }
        return bottomTopPrices
    }

    fun maxProfit(prices: IntArray, k: Int): Int {
        val bottomTopPrices = findBottomTopPrices(prices)
        val profits = computeProfits(bottomTopPrices)
        val k = minOf(k, profits.size)
        profits.sortDescending()
        var result = 0
        for (i in 0 until k) {
            result += profits[i]
        }
        return result
    }
}

The Kotlin details worth a closer look:

A pair as an IntArray. intArrayOf(bottom, top) is a lightweight pair, read back with [0] and [1]. In findBottomTopPrices the arrays hold prices, not days.
ArrayDeque as a stack. Kotlin's ArrayDeque offers addLast, last() and removeLast(). last() reads the top without removing it, and removeLast() removes it.
Short-circuit &&. In stack.isNotEmpty() && bottom < stack.last()[0], the second operand runs only if the first is true, so last() never runs on an empty stack, where it would throw. The same trick protects prices[bottom + 1] from going out of range.
Shadowing k. val k = minOf(k, profits.size) declares a new local variable with the same name as the parameter. On the right-hand side, k still means the parameter. The compiler reports a shadowing warning but accepts the code.
sortDescending() sorts a MutableList in place, which is why computeProfits returns a MutableList<Int> and not a plain List. A read-only List has no such method.
++bottom as a statement increments the variable, the same as bottom++ here. The semicolon in val n = prices.size; var i = 0 only separates two statements on one line.

Complexity

Time: O(n log n). Finding the bottoms and tops visits each day once. In computeProfits, each pair is pushed once and popped once, so the stack work is linear in the number of pairs. Sorting the candidate profits dominates at O(n log n).

Space: O(n). The pair list, the stack and the profit list each hold at most about one entry per pair.

Edge cases and common mistakes

A falling chart. [7, 6, 4, 3, 1] produces a single pair whose bottom is the last price and top is the same price, a profit of 0. The answer is 0.
An empty array. There are no pairs and no profits, so the result is 0.
A flat chart. Equal prices count as part of both the dip and the rise, which is why the comparisons use >= and <=. A run of equal prices doesn't produce extra pairs.
k = 0 or a very large k. [1, 5, 1, 5] with k = 0 gives 0 and with k = 10 gives 8. The minOf cap keeps the loop from reading past the end of profits.
Taking the largest raw trades. Sorting the unmerged pairs by profit and summing the top k fails on charts like [1, 5, 2, 6] with k = 1, where the best answer, 5, comes from merging both pairs.
Using List where a MutableList is needed. Calling sortDescending() on a read-only List doesn't compile.

Testing

[2, 4, 1] with k = 2 gives 2.
[3, 2, 6, 5, 0, 3] with k = 2 gives 7, and with k = 1 gives 4.
[1, 2, 3, 4, 5] with k = 2 gives 4.
[1, 5, 1, 5] with k = 10 gives 8.
[7, 6, 4, 3, 1] with k = 3 gives 0.
[] with k = 3 gives 0.

Keep going with KotlinAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free