Start learning free

LESSON 350 OF 414 · DYNAMIC PROGRAMMING

Best Time to Buy and Sell Stock in Scala: O(k * n) with a Running Best

Practice O(n2) 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 8
    Ctrl + Enter to check

    SCALA · DYNAMIC PROGRAMMING

    Best Time To Buy Stock in Scala, explained

    Problem overview

    The method maxProfit(prices, k) receives the price of a stock on each day and a limit k on the number of transactions. A transaction is a purchase followed by a sale, and you can hold only one share at a time. The method returns the largest total profit, and 0 when no trade makes money.

    For prices = [3, 2, 6, 5, 0, 3] and 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. For prices = [2, 4, 1] and k = 2 the answer is 2, from a single trade.

    Approach: remove the inner loop

    The table is the same as in the slower solution to this problem. dp(t)(i) is the best profit using at most t transactions in the first i days. That solution tries every buy day j for every cell, which costs O(k * n^2) in total.

    The inner loop is a search for the best buy day. Selling on day i gives prices(i - 1) plus the best value of "profit before buying, minus the price paid", taken over the earlier days. That second part doesn't depend on i at all, apart from which days have been seen so far. So it can be kept as a running maximum:

    maxDiff = the largest dp(t - 1)(j) - prices(j - 1) over the days j seen so far

    Each cell then costs a constant amount of work: dp(t)(i) is the best of three options.

    • dp(t)(i - 1): no sale on day i.
    • dp(t - 1)(i): use one transaction fewer.
    • maxDiff + prices(i - 1): sell on day i a share bought on the best earlier day.

    The running value is updated after the cell is computed, with day i's own candidate dp(t - 1)(i) - prices(i - 1). That order matters: a share bought on day i can only be sold on a later day.

    Here is row t = 2 for prices = [3, 2, 6, 5, 0, 3]. Row 1 is 0, 0, 0, 4, 4, 4, 4 for 0 to 6 days. In row 2, maxDiff stays between -3 and -1 through day 4, and the cells for days 1 to 4 are 0, 0, 4, 4. On day 5 the cell is still 4, but then maxDiff becomes dp(1)(5) - 0 = 4: row 1 says 4 can be earned in the first five days, and the share costs 0. On day 6 the sale at 3 gives 4 + 3 = 7.

    Step by step

    1. Store prices.length in n.
    2. Create dp as an Array.ofDim[Int](k + 1, n + 1) of zeros. Row 0 and column 0 stay 0.
    3. Loop over the transaction count t from 1 to k.
    4. Declare maxDiff with var, inside that loop, so each row starts fresh. Its starting value is Int.MinValue: no day has been seen yet, so there is no valid buy, and any real candidate must win against it.
    5. Loop over the days i from 1 to n.
    6. Compute dp(t)(i) as the best of the three options, nesting one math.max inside another.
    7. Update maxDiff with dp(t - 1)(i) - prices(i - 1), keeping the larger value.
    8. Return dp(k)(n).

    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) {
            var maxDiff = Int.MinValue
            for (i <- 1 to n) {
                dp(t)(i) = math.max(math.max(dp(t)(i - 1), dp(t - 1)(i)), maxDiff + prices(i - 1))
                maxDiff = math.max(maxDiff, dp(t - 1)(i) - prices(i - 1))
            }
        }
        return dp(k)(n)
    }
}

    Some Scala details in this code:

    • var against val. maxDiff changes on every day, so it must be a var. dp is a val: the name is fixed, though the cells inside can still be assigned.
    • Int.MinValue is -2147483648, the smallest Int. It makes a good "nothing yet" starting value for a maximum. It is risky to add to, because Int overflow wraps silently in Scala. Here the first addition is Int.MinValue + prices(0), and a non-negative price only moves the value up, so it can't wrap. A negative price would wrap to a huge positive number.
    • Why not 0 as the starting value? The difference dp(t - 1)(i) - prices(i - 1) is often negative. A start of 0 would act as a free purchase, and the method would report profits that no sequence of trades can reach.
    • Nested math.max. The function takes two arguments, so three values need two calls.
    • Buying on the day the earlier trade ends. The running value uses dp(t - 1)(i), which includes day i itself. That lets one trade end and the next start on the same day. Selling and rebuying at the same price gains nothing, so the answers are the same as the slower solution's.

    Complexity

    Time: O(k * n). Two loops, with constant work per cell. The slower solution takes O(k * n^2).

    Space: O(k * n). The table has (k + 1) * (n + 1) cells. Row t uses only row t - 1, so the memory could be reduced to two rows.

    Edge cases and common mistakes

    • Updating maxDiff before computing the cell. The candidate would then be a sale on the day of the purchase, which only adds dp(t - 1)(i), a value the maximum already contains. The result would not change, but the order shows the rule that a share is sold on a later day, so keep it.
    • Declaring maxDiff outside the loop over t. The value from row t - 1 would leak into row t and mix up the rows.
    • Falling prices. [5, 4, 3, 2, 1] with k = 2 returns 0, because every sale candidate loses money and the first two options in the maximum are 0.
    • Large k. [1, 2, 3, 4, 5] with k = 100 returns 4. The extra transactions are unused, but the table still has 101 rows.
    • Empty and single-day prices. With prices = [] or one price, no sale is possible, and the method returns 0. With k = 0 the loops never run and it also returns 0.
    • Index off by one. The price on day i is prices(i - 1).

    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 = [1, 2, 3, 4, 5], k = 100 gives 4.
    • prices = [5, 4, 3, 2, 1], 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