LeetCode 309. Best Time to Buy and Sell Stock with Cooldown (Dynamic Programming)

The first step in understanding this problem is understanding the decisions we can make based on the option we have, which will allow us to form a decision tree. So, starting from index 0 we have the option to buy the stock at index/day 0 or just wait and go to the next day. If we choose to buy the stock the next day our options change to sell or wait. If we then sell, we would have to skip an index and go to the current index + 2, since there is a 1 day cool down. We pretty much repeat these decisions through recursion, until we reach an index that is out of bound and returns 0. At every recursive call, we find what decision gives us the most profit. Keep in mind that if we sell a stock we gain profit, so we add the value of the stock on the day to our profit, and if we buy a stock we lose profit so we subtract. We can combine this idea with memoization/ caching of subproblems to get a time complexity of O(2n) which simplifies to O(n) because for each day we can have buy or sell options. This will still take O(n) space as well, but we can reduce this to constant space O(1) with the bottom-up solution.

Recursive Approach

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        # T: O(n) S: O(n)
        memo = {}
        def mP(i, opt):
            if i >= len(prices):
                return 0
            if (i, opt) not in memo:
                if opt == 'b':
                    res = mP(i+1, 's') - prices[i]
                else:
                    res = mP(i+2, 'b') + prices[i]
                memo[(i, opt)] = max(mP(i+1, opt), res)
            return memo[(i, opt)]
        return mP(0, 'b')

Iterative Approach

To get the iterative solution, we can simply create a 2d matrix using the sell/buy options and every day/index, and we add extra two indices for the possible out-of-bound. We can fill out this table index-by-index in reverse, using the decisions we make in the recursive approach. The below solution still takes O(n) space.

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        # T: O(n) S: O(n)
        n = len(prices)
        dp = [[0]*(n+2) for _ in range(2)]
        for i in reversed(range(n)):
            dp[0][i] = max(dp[0][i+1], dp[1][i+1]-prices[i]) # buy
            dp[1][i] = max(dp[1][i+1], dp[0][i+2]+prices[i]) # sell
        return dp[0][0]

We can reduce the space complexity to O(1) by realizing that we only look back at the previous buy, the previous previous buy, and the previous sell. So we only store those values. This is the most optimized solution I came up with O(n) time complexity and O(1) space complexity. 

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        # T: O(n) S: O(1)
        prevPrevBuy, prevBuy = 0, 0
        prevSell = 0
        for i in reversed(range(len(prices))):
            curBuy = max(prevBuy, prevSell-prices[i])
            curSell = max(prevSell, prevPrevBuy+prices[i])
            prevPrevBuy, prevBuy = prevBuy, curBuy
            prevSell = curSell
        return prevBuy

It’s really satisfying to see now how the end result is so simple, but the process is pretty complex!