Topics › Dynamic Programming › DP on Stocks

Best Time to Buy and Sell Stock with Cooldown

Medium TUF ID 56 A2Z

Problem

prices[i] is the price of a stock on day i. You may make as many transactions as you like, holding at most one share at a time, with one extra rule: after you sell, you cannot buy on the next day (a one-day cooldown). Return the maximum profit.

Constraints: 1 ≤ n ≤ 5000, 0 ≤ prices[i] ≤ 1000.

Example 1

prices = [1, 2, 3, 0, 2]
answer = 3          // buy 1, sell 2, cooldown, buy 0, sell 2: (2 - 1) + (2 - 0)

Example 2

prices = [1, 3, 1, 3]
answer = 2          // buy 1, sell 3; day 2 is a cooldown, so the second rise is lost

Without the cooldown, Example 2 would give 4 (two trades of +2). The cooldown means that splitting one long trade into two short ones now has a hidden cost: a day on which you cannot buy. That is why the “sum all rises” greedy from Stock II no longer works, and we need the DP.

Approach 1: Recursion

Intuition. Use the same state as Stock II: f(i, canBuy) is the best profit from day i onward, where canBuy = 1 means no share is held. The cooldown changes exactly one transition. When we sell on day i, day i + 1 is forbidden for buying. The cleanest way to encode that is to skip day i + 1 entirely and continue from f(i + 2, 1). Skipping is safe: on the cooldown day we hold nothing and cannot buy, so there is nothing else we could do on it.

The answer is f(0, 1).

Algorithm.

  1. If i ≥ n, return 0.
  2. Apply the buy or sell recurrence; a sell continues from i + 2.
static int recursive(int[] prices) {
    return solve(0, 1, prices);
}
static int solve(int i, int canBuy, int[] prices) {
    if (i >= prices.length) return 0;
    if (canBuy == 1) {
        return Math.max(-prices[i] + solve(i + 1, 0, prices),     // buy today
                        solve(i + 1, 1, prices));                 // skip
    }
    return Math.max(prices[i] + solve(i + 2, 1, prices),          // sell today, tomorrow is the cooldown
                    solve(i + 1, 0, prices));                     // keep holding
}

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

Approach 2: Memoization

Intuition. The state space is still 2n. The jump to i + 2 does not create new states, it only changes which state a sell leads to. Cache each (i, canBuy) the first time it is solved; -1 is a safe marker because profits are never negative.

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

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

Approach 3: Tabulation

Intuition. Row i now reads rows i + 1 and i + 2. Allocate n + 2 rows so that both dp[n] and dp[n + 1] exist and are zero, then fill from i = n - 1 downward. No special case is needed for the last day, because dp[n + 1] absorbs the jump.

Algorithm.

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

Dry run on prices = [1, 2, 3, 0, 2] (dp[i][1] = free to buy, dp[i][0] = holding):

i price dp[i][1] dp[i][0]
6, 5 - 0 0
4 2 max(-2 + 0, 0) = 0 max(2 + 0, 0) = 2
3 0 max(0 + 2, 0) = 2 max(0 + 0, 2) = 2
2 3 max(-3 + 2, 2) = 2 max(3 + 0, 2) = 3
1 2 max(-2 + 3, 2) = 2 max(2 + 2, 3) = 4
0 1 max(-1 + 4, 2) = 3 max(1 + 2, 4) = 4

Look at dp[1][0] = 4: holding a share on day 1, selling at 2 jumps to dp[3][1] = 2 (buy 0, sell 2), total 4. Then dp[0][1] = -1 + 4 = 3 is the answer.

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

Approach 4: Space optimised

Intuition. Day i needs three numbers from the future: dp[i + 1][1], dp[i + 1][0] and dp[i + 2][1]. Keep them in buy1, hold1 and buy2. After computing day i, shift the window: the old buy1 becomes the new buy2, and the freshly computed values become buy1 and hold1. Note that dp[i + 2][0] is never read, so it does not need a variable.

static int spaceOptimized(int[] prices) {
    int buy1 = 0, hold1 = 0;                            // dp[i + 1][1], dp[i + 1][0]
    int buy2 = 0;                                       // dp[i + 2][1]
    for (int i = prices.length - 1; i >= 0; i--) {
        int buy = Math.max(-prices[i] + hold1, buy1);
        int hold = Math.max(prices[i] + buy2, hold1);
        buy2 = buy1;
        buy1 = buy;
        hold1 = hold;
    }
    return buy1;
}

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

Why the cooldown is enforced exactly. Every sell transition jumps over one day in the free state, so a buy can never happen on the day right after a sell. Conversely, any legal plan has, after each sell, a day on which it does nothing, so it corresponds to a path through these transitions. The DP therefore searches precisely the set of legal plans.

Edge cases and pitfalls