Topics › Dynamic Programming › DP on Stocks

Best time to buy and sell stock III

Medium TUF ID 303 A2Z

Problem

prices[i] is the price of a stock on day i. You may complete at most two transactions, and you may not hold more than one share at a time (sell before buying again). Return the maximum profit.

Constraints: 1 ≤ n ≤ 10^5, 0 ≤ prices[i] ≤ 10^5.

Example 1

prices = [3, 3, 5, 0, 0, 3, 1, 4]
answer = 6          // buy 0, sell 3 (+3); buy 1, sell 4 (+3)

Example 2

prices = [1, 2, 4, 2, 5, 7, 2, 4, 9, 0]
answer = 13         // buy 1, sell 7 (+6); buy 2, sell 9 (+7)

In Example 2 there are three separate rises (1→4, 2→7, 2→9). With unlimited trades the answer would be 15, but with only two we must merge the first two rises into one trade 1→7. A greedy rule cannot decide which rises to merge; the DP below does it by tracking how many transactions are left.

Approach 1: Recursion

Intuition. Start from the Stock II state (i, canBuy) and add one more piece of information the future depends on: cap, the number of transactions still allowed. A transaction is counted as finished when we sell, so cap decreases on the sell transition only. When cap reaches 0, nothing more can be earned.

f(i, canBuy, cap) = best profit from day i onward:

The answer is f(0, 1, 2). Counting on the sell rather than the buy is a choice; counting on the buy also works, but then the base case must still allow selling a share bought with the last transaction.

Algorithm.

  1. If i == n or cap == 0, return 0.
  2. Apply the buy or sell recurrence according to canBuy.
static int recursive(int[] prices) {
    return solve(0, 1, 2, prices);
}
static int solve(int i, int canBuy, int cap, int[] prices) {
    if (i == prices.length || cap == 0) return 0;
    if (canBuy == 1) {
        return Math.max(-prices[i] + solve(i + 1, 0, cap, prices),       // buy today
                        solve(i + 1, 1, cap, prices));                   // skip
    }
    return Math.max(prices[i] + solve(i + 1, 1, cap - 1, prices),        // sell: one transaction done
                    solve(i + 1, 0, cap, prices));                       // keep holding
}

Complexity. Time O(2^n), space O(n) stack.

Approach 2: Memoization

Intuition. The state space is n × 2 × 3 (cap is 0, 1 or 2), only 6n states. Cache them in a 3D array initialised to -1; profits are never negative, so -1 cannot be confused with a real answer.

static int memoization(int[] prices) {
    int[][][] dp = new int[prices.length][2][3];
    for (int[][] a : dp) for (int[] b : a) Arrays.fill(b, -1);
    return memo(0, 1, 2, prices, dp);
}
static int memo(int i, int canBuy, int cap, int[] prices, int[][][] dp) {
    if (i == prices.length || cap == 0) return 0;
    if (dp[i][canBuy][cap] != -1) return dp[i][canBuy][cap];
    int best;
    if (canBuy == 1) {
        best = Math.max(-prices[i] + memo(i + 1, 0, cap, prices, dp), memo(i + 1, 1, cap, prices, dp));
    } else {
        best = Math.max(prices[i] + memo(i + 1, 1, cap - 1, prices, dp), memo(i + 1, 0, cap, prices, dp));
    }
    return dp[i][canBuy][cap] = best;
}

Complexity. Time O(n) (6n states, O(1) each). Space O(n) for the table plus O(n) stack.

Approach 3: Tabulation

Intuition. Day i reads only day i + 1, so iterate i from n - 1 down to 0. The base cases are free: row n is all zeros, and cap = 0 is never written, so it stays zero. Inside a day, the order of cap and canBuy does not matter, because every value read comes from row i + 1.

Algorithm.

  1. dp[n + 1][2][3], all zero.
  2. For i from n - 1 down to 0, and cap from 1 to 2, fill both canBuy values.
  3. Return dp[0][1][2].
static int tabulation(int[] prices) {
    int n = prices.length;
    int[][][] dp = new int[n + 1][2][3];                // dp[n][*][*] = 0 and dp[*][*][0] = 0
    for (int i = n - 1; i >= 0; i--) {
        for (int cap = 1; cap <= 2; cap++) {
            dp[i][1][cap] = Math.max(-prices[i] + dp[i + 1][0][cap], dp[i + 1][1][cap]);
            dp[i][0][cap] = Math.max(prices[i] + dp[i + 1][1][cap - 1], dp[i + 1][0][cap]);
        }
    }
    return dp[0][1][2];
}

Complexity. Time O(n), space O(n).

Approach 4: Space optimised

Intuition. Keep two small 2 × 3 tables: ahead for day i + 1 and cur for day i, and swap them after each day. Swapping references instead of copying keeps the loop allocation-free.

static int spaceOptimized(int[] prices) {
    int[][] ahead = new int[2][3], cur = new int[2][3];
    for (int i = prices.length - 1; i >= 0; i--) {
        for (int cap = 1; cap <= 2; cap++) {
            cur[1][cap] = Math.max(-prices[i] + ahead[0][cap], ahead[1][cap]);
            cur[0][cap] = Math.max(prices[i] + ahead[1][cap - 1], ahead[0][cap]);
        }
        int[][] t = ahead; ahead = cur; cur = t;
    }
    return ahead[1][2];
}

Complexity. Time O(n), space O(1) (twelve integers).

Approach 5: Four running states

Intuition. The same DP can be run forwards with an even simpler state: at the end of each day, what is the best balance (money in hand, starting from 0) after reaching each milestone?

Each milestone is reached either earlier (keep the old value) or today, by acting from the previous milestone: buy1 = max(buy1, -p), sell1 = max(sell1, buy1 + p), buy2 = max(buy2, sell1 - p), sell2 = max(sell2, buy2 + p). Updating in this order lets a milestone use a same-day value from the previous milestone, which corresponds to buying and selling on the same day: profit 0, so it never creates fake profit. sell2 already covers plans with fewer trades, because an unused transaction can be a same-day buy and sell.

static int fourStates(int[] prices) {
    int buy1 = Integer.MIN_VALUE, sell1 = 0;            // best balance after the first buy / first sell
    int buy2 = Integer.MIN_VALUE, sell2 = 0;            // best balance after the second buy / second sell
    for (int p : prices) {
        buy1 = Math.max(buy1, -p);
        sell1 = Math.max(sell1, buy1 + p);
        buy2 = Math.max(buy2, sell1 - p);
        sell2 = Math.max(sell2, buy2 + p);
    }
    return sell2;
}

Dry run on prices = [3, 3, 5, 0, 0, 3, 1, 4]:

p buy1 sell1 buy2 sell2
3 -3 0 -3 0
3 -3 0 -3 0
5 -3 2 -3 2
0 0 2 2 2
0 0 2 2 2
3 0 3 2 5
1 0 3 2 5
4 0 4 2 6

On the fourth day buy2 = 2: we sold once for +2 (buy 3, sell 5) and bought again at 0. The final sell2 = 6 comes from buy2 = 2 plus selling at 4, which is the plan “buy 3, sell 5, buy 0, sell 4”, also worth 6.

Complexity. Time O(n), space O(1). This is the form most people write in interviews, but it is worth knowing that it is exactly the DP of Approach 4 read in the opposite direction.

Edge cases and pitfalls