Start learning free

LESSON 349 OF 414 · DYNAMIC PROGRAMMING

Best Time to Buy and Sell Stock in Rust: DP With 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

    RUST · DYNAMIC PROGRAMMING

    Best Time To Buy Stock in Rust, explained

    Problem overview

    The function max_profit receives prices, where prices[d] is a stock's price on a given day, and k, the maximum number of transactions allowed. A transaction is one buy followed by one sell. You can't hold two purchases at once: the stock must be sold before you buy again. The function returns the largest total profit.

    In the second example, prices = [3, 2, 6, 5, 0, 3] and k = 2. Buying at 2 and selling at 6 earns 4, and buying at 0 and selling at 3 earns 3. The total is 7. In the first example, [2, 4, 1], one transaction (buy at 2, sell at 4) earns 2, and a second one can't add anything.

    Approach: best profit for each number of transactions and days

    With one transaction the task is to pick the lowest earlier price for each sale. With several transactions, the choices interact, so a table keeps them organised. Let dp[t][i] be the best profit using at most t transactions during the first i days. Row 0 (no transactions) and column 0 (no days) are always 0.

    To compute dp[t][i], consider what happens on day i:

    • No sale on day i. The answer is no worse than dp[t][i - 1], the same budget with one day less. It is also no worse than dp[t - 1][i], because "at most" t transactions includes using fewer.
    • A sale on day i. Choose the day j before i on which that stock was bought. The profit of the last transaction is prices[i] - prices[j], and before day j you used at most t - 1 transactions, worth dp[t - 1][j - 1]. Try every j and keep the best.

    This is the straightforward version of the idea: for every cell it scans all possible buy days. A faster version that removes that scan exists, and the next lesson on this problem covers it.

    Step by step

    1. Store the length of prices in n, because it is used constantly, and rebind k as a usize so it can size the table.
    2. Create the table. dp is a mutable vector of k + 1 rows, each with n + 1 zeros of type i32.
    3. Loop over t from 1 to k. Row 0 stays 0, so the loop skips it.
    4. Loop over i from 1 to n. Column 0 also stays 0.
    5. Set the starting value for the "no sale on day i" case: the larger of dp[t][i - 1] and dp[t - 1][i], using max.
    6. Loop over the buy day j from 1 to i - 1. Buying has to happen before selling.
    7. Compute the profit of buying on day j and selling on day i. Days are numbered from 1 here but prices is indexed from 0, so the prices are prices[j - 1] and prices[i - 1].
    8. Update the cell if this last transaction, added to the best profit before day j with one transaction fewer, beats what is stored.
    9. Return dp[k][n] as the last expression.

    The solution in Rust

    impl Solution {
    fn max_profit(prices: Vec<i32>, k: i32) -> i32 {
        let n = prices.len();
        let k = k as usize;
        let mut dp = vec![vec![0i32; n + 1]; k + 1];
        for t in 1..=k {
            for i in 1..=n {
                dp[t][i] = dp[t][i - 1].max(dp[t - 1][i]);
                for j in 1..i {
                    let profit = prices[i - 1] - prices[j - 1];
                    dp[t][i] = dp[t][i].max(dp[t - 1][j - 1] + profit);
                }
            }
        }
        dp[k][n]
    }
}

    The Rust details worth knowing:

    • Methods on numbers. a.max(b) returns the larger of two integers. It is a method of the number type, so dp[t][i - 1].max(dp[t - 1][i]) reads left to right without importing anything.
    • Copying out of the vector. Indexing a Vec<i32> gives an i32, which is Copy. That is why dp[t][i] = dp[t][i].max(...) reads and writes the same cell without any borrowing trouble.
    • Ranges. 1..=k and 1..=n include their upper end. 1..i excludes i, so j stops at i - 1, exactly the days before the sale.
    • Two sets of indexes. The table counts days from 1 and prices counts from 0. The - 1 in prices[i - 1] converts between them, and dp[t - 1][j - 1] means "the first j - 1 days".
    • Row cloning. The nested vec! makes k + 1 independent rows. The literal 0i32 makes the cells the same type as the prices.
    • Overflow. The profit is a difference of two prices, and sums of such differences stay within an i32 for realistic price lists. In the optimized builds the lessons use, an overflow would wrap silently rather than stop the program.

    Complexity

    Time: O(k * n²). The loops over t, i and j are nested. The inner loop runs up to n times for each of the k * n cells.

    Space: O(k * n). The table holds (k + 1) * (n + 1) numbers.

    The cubic cost is what the faster lesson improves: it keeps a running best value for the buy day, so the inner loop disappears and the time drops to O(k * n).

    Edge cases and common mistakes

    • Prices that only fall. [5, 4, 3] with k = 2 gives 0. Every profit is negative, and the starting value from the no-sale case keeps the cell at 0.
    • No transactions. With k = 0 the loops over t don't run and dp[0][n] is 0.
    • Empty or single-day input. With no prices or one price, no sale is possible and the result is 0. The inner loop is empty, and n = 0 reads only dp[k][0].
    • More transactions than needed. For [7, 1, 5, 3, 6, 4] and k = 1 the answer is 5 (buy at 1, sell at 6). Extra transactions never reduce the result, because "at most" allows using fewer.
    • Ignoring the "at most" rule. Without the dp[t - 1][i] term, a budget of t would be forced to use all t transactions.
    • Reusing a day. The previous profit comes from dp[t - 1][j - 1], the days before the buy day. Using dp[t - 1][j] would let one day serve as both a sale and a purchase.
    • Index off by one. Writing prices[i] instead of prices[i - 1] reads past the end on day n and panics.

    Testing

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