Start learning free

LESSON 349 OF 414 · DYNAMIC PROGRAMMING

Best Time to Buy and Sell Stock in Scala: At Most K Transactions

Practice O(n3) Read the full explanation

Say you have an array prices 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 (i.e. you must sell the stock before you buy again).

Example 1:

Input: prices = [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: prices = [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 9
    Ctrl + Enter to check

    SCALA · DYNAMIC PROGRAMMING

    Best Time To Buy Stock in Scala, explained

    Problem overview

    The method maxProfit(prices, k) receives an array where prices(i) is a stock's price on day i + 1, and a limit k on the number of transactions. A transaction is one purchase followed by one sale. You may hold at most one share at a time, so you must sell before you buy again. The method returns the largest total profit, and 0 if no trade makes money.

    With 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 prices = [2, 4, 1] and k = 2 the answer is 2, from one trade. The second transaction would lose money, and "at most" lets you skip it.

    Approach: best profit for each number of trades and days

    Choosing the trades greedily fails, because a good first trade can block a better pair of trades. Dynamic programming looks at all the possibilities in a table.

    Let dp(t)(i) be the best profit using at most t transactions during the first i days. There is a row for each number of transactions (0 to k) and a column for each number of days (0 to n). Row 0 and column 0 stay 0: no trades or no days means no profit.

    To fill dp(t)(i), consider what happens on day i:

    • No sale on day i. The best profit is the same as for one day fewer, dp(t)(i - 1). It is also at least as good as using one transaction fewer, dp(t - 1)(i), since "at most" allows unused transactions.
    • A sale on day i. The share was bought on some earlier day j. This last trade earns prices(i - 1) - prices(j - 1), and the earlier trades happened in the days before j with at most t - 1 transactions: dp(t - 1)(j - 1).

    The entry is the largest value among these options, trying every j from 1 to i - 1.

    The days are numbered from 1 in the table, but the array is indexed from 0, so day d's price is prices(d - 1). This is the source of most mistakes in the code.

    Step by step

    1. Store prices.length in n, since the algorithm uses it many times.
    2. Create dp with Array.ofDim[Int](k + 1, n + 1): k + 1 rows, n + 1 columns, all zeros.
    3. Loop over the transaction count t from 1 to k. Row 0 stays 0.
    4. Loop over the days i from 1 to n. Column 0 stays 0.
    5. Start with the "no sale" value: math.max(dp(t)(i - 1), dp(t - 1)(i)). This is only a starting value, and the next loop may improve it.
    6. Loop over the buy day j from 1 to i - 1. The purchase must come before the sale.
    7. Compute the profit of buying on day j and selling on day i: prices(i - 1) - prices(j - 1).
    8. Combine it with the earlier trades: if dp(t - 1)(j - 1) + profit beats the current dp(t)(i), keep it.
    9. Return dp(k)(n) as the last expression.

    The solution in Scala

    class Solution {
        def maxProfit(prices: Array[Int], k: Int): Int = {
            val n = prices.length
            val dp = Array.ofDim[Int](k + 1, n + 1)
            for (t <- 1 to k) {
                for (i <- 1 to n) {
                    dp(t)(i) = math.max(dp(t)(i - 1), dp(t - 1)(i))
                    for (j <- 1 until i) {
                        val profit = prices(i - 1) - prices(j - 1)
                        dp(t)(i) = math.max(dp(t)(i), dp(t - 1)(j - 1) + profit)
                    }
                }
            }
            dp(k)(n)
        }
    }
    

    Some Scala details in this code:

    • Array.ofDim[Int](k + 1, n + 1) creates the whole 2D array with zeros, and dp(t)(i) reads and writes its cells. Array indexing uses parentheses in Scala.
    • 1 to k against 1 until i. to includes its upper bound and until excludes it. The buy day j must stay strictly below i, so it uses until. When i is 1 the range 1 until 1 is empty and the loop body doesn't run.
    • math.max compares two values. To compare three, nest one call in another.
    • The array is a val, but its cells change. val dp fixes which array the name refers to, not what is inside.
    • The last expression is the result. dp(k)(n) is returned without return.
    • Negative profit is harmless. A losing trade produces a smaller candidate than the starting value, so math.max ignores it. This is how "at most" is honoured.

    Complexity

    Time: O(k * n^2). For each of the k rows and each of the n days, the inner loop tries up to n buy days. A faster variant keeps a running best of "profit so far minus the buy price" and removes the inner loop, which gives O(k * n).

    Space: O(k * n). The table has (k + 1) * (n + 1) cells. Since a very large k allocates a large table, a common refinement is to cap it at n / 2: more trades than that can't all be profitable.

    Edge cases and common mistakes

    • No profitable trade. For falling prices such as [5, 4, 3, 2, 1] every candidate is negative, and the method returns 0.
    • k = 0. The loops don't run, and dp(0)(n) is 0.
    • Empty or one-day prices. With no days, or one day, there is nothing to sell after buying, and the answer is 0.
    • More transactions than needed. For [3, 2, 6, 5, 0, 3], k = 3 still returns 7, since only two rises are worth trading.
    • Using the wrong earlier row. The trades before the buy must use at most t - 1 transactions: dp(t - 1)(j - 1). Reading row t would reuse the transaction being decided.
    • Off-by-one on prices. Day i is at index i - 1 in the array. Writing prices(i) reads the next day and fails with an exception on the last day.
    • Overlapping trades. The buy for the last trade happens after the days that dp(t - 1)(j - 1) covers, so there are never two shares at once.

    Testing

    • prices = [2, 4, 1], k = 2 gives 2.
    • prices = [3, 2, 6, 5, 0, 3], k = 2 gives 7.
    • prices = [3, 2, 6, 5, 0, 3], k = 1 gives 4.
    • prices = [5, 4, 3, 2, 1], k = 2 gives 0.
    • prices = [7], k = 2 gives 0.
    Keep going with ScalaAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free