Skip to content

309. Best Time to Buy and Sell Stock with Cooldown

DifficultyMedium

Description

You are given an array prices where prices[i] is the price of a given stock on the ith day.

Find the maximum profit you can achieve. You may complete as many transactions as you like (i.e., buy one and sell one share of the stock multiple times) with the following restrictions:

  • After you sell your stock, you cannot buy stock on the next day (i.e., cooldown one day).

Note: You may not engage in multiple transactions simultaneously (i.e., you must sell the stock before you buy again).

 

Example 1:

Input: prices = [1,2,3,0,2]
Output: 3
Explanation: transactions = [buy, sell, cooldown, buy, sell]

Example 2:

Input: prices = [1]
Output: 0

 

Constraints:

  • 1 <= prices.length <= 5000
  • 0 <= prices[i] <= 1000

Solutions

Thinking

We may trade many times, but a sell forces a one-day cooldown. A raw decision tree over day, holding, and cooldown revisits the same states.

Compress to \((i,j)\): starting at day \(i\), whether we hold. Skip the day; if holding, sell and jump to \(i+2\); if free, buy and start holding. Memoization evaluates each state once; the extra day after a sell is the cooldown.

We design a function \(dfs(i, j)\), which represents the maximum profit that can be obtained starting from the \(i\)th day with state \(j\). The values of \(j\) are \(0\) and \(1\), respectively representing currently not holding a stock and holding a stock. The answer is \(dfs(0, 0)\).

The execution logic of the function \(dfs(i, j)\) is as follows:

If \(i \geq n\), it means that there are no more stocks to trade, so return \(0\);

Otherwise, we can choose not to trade, then \(dfs(i, j) = dfs(i + 1, j)\). We can also trade stocks. If \(j > 0\), it means that we currently hold a stock and can sell it, then \(dfs(i, j) = prices[i] + dfs(i + 2, 0)\). If \(j = 0\), it means that we currently do not hold a stock and can buy, then \(dfs(i, j) = -prices[i] + dfs(i + 1, 1)\). Take the maximum value as the return value of the function \(dfs(i, j)\).

The answer is \(dfs(0, 0)\).

To avoid repeated calculations, we use the method of memoization search, and use an array \(f\) to record the return value of \(dfs(i, j)\). If \(f[i][j]\) is not \(-1\), it means that it has been calculated, and we can directly return \(f[i][j]\).

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the length of the array \(prices\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        @cache
        def dfs(i: int, j: int) -> int:
            if i >= len(prices):
                return 0
            ans = dfs(i + 1, j)
            if j:
                ans = max(ans, prices[i] + dfs(i + 2, 0))
            else:
                ans = max(ans, -prices[i] + dfs(i + 1, 1))
            return ans

        return dfs(0, 0)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
    private int[] prices;
    private Integer[][] f;

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

    private int dfs(int i, int j) {
        if (i >= prices.length) {
            return 0;
        }
        if (f[i][j] != null) {
            return f[i][j];
        }
        int ans = dfs(i + 1, j);
        if (j > 0) {
            ans = Math.max(ans, prices[i] + dfs(i + 2, 0));
        } else {
            ans = Math.max(ans, -prices[i] + dfs(i + 1, 1));
        }
        return f[i][j] = ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int n = prices.size();
        int f[n][2];
        memset(f, -1, sizeof(f));
        function<int(int, int)> dfs = [&](int i, int j) {
            if (i >= n) {
                return 0;
            }
            if (f[i][j] != -1) {
                return f[i][j];
            }
            int ans = dfs(i + 1, j);
            if (j) {
                ans = max(ans, prices[i] + dfs(i + 2, 0));
            } else {
                ans = max(ans, -prices[i] + dfs(i + 1, 1));
            }
            return f[i][j] = ans;
        };
        return dfs(0, 0);
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
func maxProfit(prices []int) int {
    n := len(prices)
    f := make([][2]int, n)
    for i := range f {
        f[i] = [2]int{-1, -1}
    }
    var dfs func(i, j int) int
    dfs = func(i, j int) int {
        if i >= n {
            return 0
        }
        if f[i][j] != -1 {
            return f[i][j]
        }
        ans := dfs(i+1, j)
        if j > 0 {
            ans = max(ans, prices[i]+dfs(i+2, 0))
        } else {
            ans = max(ans, -prices[i]+dfs(i+1, 1))
        }
        f[i][j] = ans
        return ans
    }
    return dfs(0, 0)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
function maxProfit(prices: number[]): number {
    const n = prices.length;
    const f: number[][] = Array.from({ length: n }, () => Array.from({ length: 2 }, () => -1));
    const dfs = (i: number, j: number): number => {
        if (i >= n) {
            return 0;
        }
        if (f[i][j] !== -1) {
            return f[i][j];
        }
        let ans = dfs(i + 1, j);
        if (j) {
            ans = Math.max(ans, prices[i] + dfs(i + 2, 0));
        } else {
            ans = Math.max(ans, -prices[i] + dfs(i + 1, 1));
        }
        return (f[i][j] = ans);
    };
    return dfs(0, 0);
}

Solution 2: Dynamic Programming

Thinking

Solution 1 fills a suffix table from the end. The same choices can be written forward: \(f[i][0/1]\) is the best profit after day \(i\), free or holding. Free comes from staying free or selling today; holding comes from staying put or buying after cooldown, i.e. from \(f[i-2][0]\).

Fill left to right; the answer is free on the last day. Time stays \(O(n)\).

We can also use dynamic programming to solve this problem.

We define \(f[i][j]\) to represent the maximum profit that can be obtained on the \(i\)th day with state \(j\). The values of \(j\) are \(0\) and \(1\), respectively representing currently not holding a stock and holding a stock. Initially, \(f[0][0] = 0\), \(f[0][1] = -prices[0]\).

When \(i \geq 1\), if we currently do not hold a stock, then \(f[i][0]\) can be obtained by transitioning from \(f[i - 1][0]\) and \(f[i - 1][1] + prices[i]\), i.e., \(f[i][0] = \max(f[i - 1][0], f[i - 1][1] + prices[i])\). If we currently hold a stock, then \(f[i][1]\) can be obtained by transitioning from \(f[i - 1][1]\) and \(f[i - 2][0] - prices[i]\), i.e., \(f[i][1] = \max(f[i - 1][1], f[i - 2][0] - prices[i])\). The final answer is \(f[n - 1][0]\).

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the length of the array \(prices\).

1
2
3
4
5
6
7
8
9
class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        n = len(prices)
        f = [[0] * 2 for _ in range(n)]
        f[0][1] = -prices[0]
        for i in range(1, n):
            f[i][0] = max(f[i - 1][0], f[i - 1][1] + prices[i])
            f[i][1] = max(f[i - 1][1], f[i - 2][0] - prices[i])
        return f[n - 1][0]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
class Solution {
    public int maxProfit(int[] prices) {
        int n = prices.length;
        int[][] f = new int[n][2];
        f[0][1] = -prices[0];
        for (int i = 1; i < n; i++) {
            f[i][0] = Math.max(f[i - 1][0], f[i - 1][1] + prices[i]);
            f[i][1] = Math.max(f[i - 1][1], (i > 1 ? f[i - 2][0] : 0) - prices[i]);
        }
        return f[n - 1][0];
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int n = prices.size();
        int f[n][2];
        memset(f, 0, sizeof(f));
        f[0][1] = -prices[0];
        for (int i = 1; i < n; ++i) {
            f[i][0] = max(f[i - 1][0], f[i - 1][1] + prices[i]);
            f[i][1] = max(f[i - 1][1], (i > 1 ? f[i - 2][0] : 0) - prices[i]);
        }
        return f[n - 1][0];
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
func maxProfit(prices []int) int {
    n := len(prices)
    f := make([][2]int, n)
    f[0][1] = -prices[0]
    for i := 1; i < n; i++ {
        f[i][0] = max(f[i-1][0], f[i-1][1]+prices[i])
        if i > 1 {
            f[i][1] = max(f[i-1][1], f[i-2][0]-prices[i])
        } else {
            f[i][1] = max(f[i-1][1], -prices[i])
        }
    }
    return f[n-1][0]
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
function maxProfit(prices: number[]): number {
    const n = prices.length;
    const f: number[][] = Array.from({ length: n }, () => Array.from({ length: 2 }, () => 0));
    f[0][1] = -prices[0];
    for (let i = 1; i < n; ++i) {
        f[i][0] = Math.max(f[i - 1][0], f[i - 1][1] + prices[i]);
        f[i][1] = Math.max(f[i - 1][1], (i > 1 ? f[i - 2][0] : 0) - prices[i]);
    }
    return f[n - 1][0];
}

Solution 3: Dynamic Programming (Space Optimization)

Thinking

Method 2 only reads \(i-1\) and \(i-2\), so the full table is unnecessary. Three rolling variables (free two days ago, free yesterday, holding yesterday) implement the same transfers in \(O(1)\) space.

The transition only needs the previous two days, so three variables are enough and the space complexity is \(O(1)\).

1
2
3
4
5
6
class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        f, f0, f1 = 0, 0, -prices[0]
        for x in prices[1:]:
            f, f0, f1 = f0, max(f0, f1 + x), max(f1, f - x)
        return f0

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
class Solution {
    public int maxProfit(int[] prices) {
        int f = 0, f0 = 0, f1 = -prices[0];
        for (int i = 1; i < prices.length; ++i) {
            int g0 = Math.max(f0, f1 + prices[i]);
            f1 = Math.max(f1, f - prices[i]);
            f = f0;
            f0 = g0;
        }
        return f0;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int f = 0, f0 = 0, f1 = -prices[0];
        for (int i = 1; i < prices.size(); ++i) {
            int g0 = max(f0, f1 + prices[i]);
            f1 = max(f1, f - prices[i]);
            f = f0;
            f0 = g0;
        }
        return f0;
    }
};

1
2
3
4
5
6
7
func maxProfit(prices []int) int {
    f, f0, f1 := 0, 0, -prices[0]
    for _, x := range prices[1:] {
        f, f0, f1 = f0, max(f0, f1+x), max(f1, f-x)
    }
    return f0
}

1
2
3
4
5
6
7
function maxProfit(prices: number[]): number {
    let [f, f0, f1] = [0, 0, -prices[0]];
    for (const x of prices.slice(1)) {
        [f, f0, f1] = [f0, Math.max(f0, f1 + x), Math.max(f1, f - x)];
    }
    return f0;
}

Solution 4: Dynamic Programming

Thinking

We may trade many times, but a sell forces a one-day cooldown. \(n \le 5000\), so a raw decision tree is far too large.

The state is the day and whether we hold. Skipping always calls the next day before it returns, so the chain has length \(n\) and overflows the stack.

Later days are known if we walk backward. Let \(f[i][j]\) be the best profit from day \(i\) with holding flag \(j\), and fill \(i\) from \(n-1\) down to \(0\). Selling lands on \(i+2\), which is the cooldown.

Let \(f[i][j]\) be the maximum profit starting from day \(i\) in state \(j\). The values of \(j\) are \(0\) and \(1\), meaning we do not hold a stock and we hold a stock. The answer is \(f[0][0]\). Days past the end contribute \(0\), so \(f[n][j] = f[n + 1][j] = 0\).

Fill \(i\) from \(n - 1\) down to \(0\). Doing nothing keeps \(f[i + 1][j]\). If \(j > 0\), we hold a stock and may sell it, earning \(prices[i] + f[i + 2][0]\). The extra day is the cooldown. If \(j = 0\), we may buy, earning \(-prices[i] + f[i + 1][1]\). Take the larger value:

\[ f[i][j] = \max(f[i + 1][j],\ prices[i] + f[i + 2][0]) \]

when \(j > 0\), and

\[ f[i][j] = \max(f[i + 1][j],\ -prices[i] + f[i + 1][1]) \]

when \(j = 0\).

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the length of the array \(prices\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        n = len(prices)
        f = [[0] * 2 for _ in range(n + 2)]
        for i in range(n - 1, -1, -1):
            for j in range(2):
                ans = f[i + 1][j]
                if j:
                    ans = max(ans, prices[i] + f[i + 2][0])
                else:
                    ans = max(ans, -prices[i] + f[i + 1][1])
                f[i][j] = ans
        return f[0][0]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution {
    public int maxProfit(int[] prices) {
        int n = prices.length;
        int[][] f = new int[n + 2][2];
        for (int i = n - 1; i >= 0; --i) {
            for (int j = 0; j < 2; ++j) {
                int ans = f[i + 1][j];
                if (j > 0) {
                    ans = Math.max(ans, prices[i] + f[i + 2][0]);
                } else {
                    ans = Math.max(ans, -prices[i] + f[i + 1][1]);
                }
                f[i][j] = ans;
            }
        }
        return f[0][0];
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int n = prices.size();
        vector<vector<int>> f(n + 2, vector<int>(2));
        for (int i = n - 1; i >= 0; --i) {
            for (int j = 0; j < 2; ++j) {
                int ans = f[i + 1][j];
                if (j) {
                    ans = max(ans, prices[i] + f[i + 2][0]);
                } else {
                    ans = max(ans, -prices[i] + f[i + 1][1]);
                }
                f[i][j] = ans;
            }
        }
        return f[0][0];
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
func maxProfit(prices []int) int {
    n := len(prices)
    f := make([][2]int, n+2)
    for i := n - 1; i >= 0; i-- {
        for j := 0; j < 2; j++ {
            ans := f[i+1][j]
            if j > 0 {
                ans = max(ans, prices[i]+f[i+2][0])
            } else {
                ans = max(ans, -prices[i]+f[i+1][1])
            }
            f[i][j] = ans
        }
    }
    return f[0][0]
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
function maxProfit(prices: number[]): number {
    const n = prices.length;
    const f: number[][] = Array.from({ length: n + 2 }, () => Array(2).fill(0));
    for (let i = n - 1; i >= 0; --i) {
        for (let j = 0; j < 2; ++j) {
            let ans = f[i + 1][j];
            if (j) {
                ans = Math.max(ans, prices[i] + f[i + 2][0]);
            } else {
                ans = Math.max(ans, -prices[i] + f[i + 1][1]);
            }
            f[i][j] = ans;
        }
    }
    return f[0][0];
}

Comments