Best Time to Buy and Sell Stock with Cooldown

NeetCode: Best Time to Buy and Sell Stock with Cooldown · LeetCode: 309. Best Time to Buy and Sell Stock with Cooldown · Difficulty: Medium

prices[i] is the price of a stock on day i. You may complete as many transactions as you like, but you cannot hold more than one share at a time, and after you sell you cannot buy on the next day (a one-day cooldown). Return the maximum profit.

Example

Text

Input:  prices = [1, 2, 3, 0, 2]    Output: 3   (buy, sell, cooldown, buy, sell)
Input:  prices = [1]                Output: 0

Brute force: try every decision

On each day the state is whether you currently hold a share. If you do not hold one, you can skip the day or buy. If you hold one, you can skip or sell, and selling jumps ahead two days to honour the cooldown. Trying every choice on every day gives up to 2^n decision sequences.

Java

class Solution {
    private int dfs(int[] prices, int i, boolean holding) {
        if (i >= prices.length) {
            return 0;
        }
        int skip = dfs(prices, i + 1, holding);
        if (holding) {
            return Math.max(skip, prices[i] + dfs(prices, i + 2, false));
        }
        return Math.max(skip, -prices[i] + dfs(prices, i + 1, true));
    }

    public int maxProfit(int[] prices) {
        return dfs(prices, 0, false);
    }
}

public class Main {
    public static void main(String[] args) {
        Solution sol = new Solution();
        System.out.println(sol.maxProfit(new int[] {1, 2, 3, 0, 2}));
        System.out.println(sol.maxProfit(new int[] {1}));
        System.out.println(sol.maxProfit(new int[] {}));
    }
}

Output

3
0
0

Time: O(2^n). Space: O(n).

Better: memoise on (day, holding)

The answer for a state depends only on the day and whether you hold a share, so there are only 2n distinct states. Cache each result the first time it is computed. The cache is an Integer[][]: a boxed Integer can hold null for "not computed", which an int cannot.

Java

class Solution {
    // memo[i][h]: best profit from day i with h = 1 if holding; null = not computed
    private int dfs(int[] prices, int i, int holding, Integer[][] memo) {
        if (i >= prices.length) {
            return 0;
        }
        if (memo[i][holding] != null) {
            return memo[i][holding];
        }
        int skip = dfs(prices, i + 1, holding, memo);
        int best;
        if (holding == 1) {
            best = Math.max(skip, prices[i] + dfs(prices, i + 2, 0, memo));
        } else {
            best = Math.max(skip, -prices[i] + dfs(prices, i + 1, 1, memo));
        }
        memo[i][holding] = best;
        return best;
    }

    public int maxProfit(int[] prices) {
        return dfs(prices, 0, 0, new Integer[prices.length][2]);
    }
}

public class Main {
    public static void main(String[] args) {
        Solution sol = new Solution();
        System.out.println(sol.maxProfit(new int[] {1, 2, 3, 0, 2}));
        System.out.println(sol.maxProfit(new int[] {1}));
        System.out.println(sol.maxProfit(new int[] {}));
    }
}

Output

3
0
0

Time: O(n). Space: O(n).

Optimal: three-state machine

Describe the end of each day with three numbers, the best profit if you are in that state:

hold: you own a share.
sold: you sold today (so tomorrow is the cooldown).
rest: you own nothing and are free to buy.

Transitions for a price p: hold' = max(hold, rest - p) (keep it, or buy from a free state), sold' = hold + p (sell what you hold) and rest' = max(rest, sold) (stay free, or come out of the cooldown). Buying only from rest is what enforces the cooldown. The answer is max(sold, rest) since ending while holding is never best. Only yesterday's values are needed, so three variables replace the whole table.

day	price	hold	sold	rest
start		-inf	-inf	0
0	1	-1	-inf	0
1	2	-1	1	0
2	3	-1	2	1
3	0	1	-1	2
4	2	1	3	2

Java

class Solution {
    public int maxProfit(int[] prices) {
        int hold = Integer.MIN_VALUE / 2; // impossible at the start; halved so adding a price cannot overflow
        int sold = Integer.MIN_VALUE / 2;
        int rest = 0;
        for (int p : prices) {
            int prevSold = sold;
            sold = hold + p;
            hold = Math.max(hold, rest - p);
            rest = Math.max(rest, prevSold);
        }
        return Math.max(sold, rest);
    }
}

public class Main {
    public static void main(String[] args) {
        Solution sol = new Solution();
        System.out.println(sol.maxProfit(new int[] {1, 2, 3, 0, 2}));
        System.out.println(sol.maxProfit(new int[] {1}));
        System.out.println(sol.maxProfit(new int[] {}));
    }
}

Output

3
0
0

Time: O(n). Space: O(1).