309. Best Time to Buy and Sell Stock with Cooldown

Hello, reader 👋🏽 ! Welcome to day 83 of the series on Problem Solving. Through this series, I aim to pick up at least one question every day and share my approach to solving it.

Today, I will be picking up LeetCode's daily challenge problem: 309. Best Time to Buy and Sell Stock with Cooldown.

🤔 Problem Statement

Given an array prices where prices[i] is the price of a given stock on the i<sup>th</sup> 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).

E.g.:

  prices = [1,2,3,0,2]
  => 3
  prices = [1,2]
  => 1

NOTE

Before we move on to this question, if you've solved Best Time to Buy and Sell Stock and Best Time to Buy and Sell Stock II then it would be easier to understand this question. This question is an extension of these two problems.
- Hence it will be easier to understand the solution if you solve the two problems.

💬 Thought Process - Brute Force

This problem is a classic problem with decisions to either include the current element or not include the current element in the answer.
We have two different states: one is the index that indicates the current element and the other is canBuy to indicate if we can buy the current stock or not.
If we can buy the current element, indicated by true for canBuy, then we have two choices: buy with the current stock price or don't buy.
- If we buy the current stock price, then we add -prices[index] to the answer and change the state of canBuy to false.
- If we don't buy for the current stock price, then we simply recursively call for buying the next price.
- The answer for the current set of state index, canBuy is going the be the max of two calls:
  - max(buyCurrentStock, dontBuyCurrentStock)
If we can sell the current stock, as indicated by false for canBuy, then we have two choices: sell for the current price or don't sell.
- If we sell the previous stock for the current price, then we add prices[index] to the answer and change the state to true.
- If we don't sell for the current stock price, then we simply recursively call for selling with the next price.
- The answer for the current set of state index, canBuy is going the be the max of two calls:
  - max(sellCurrentStock, dontSellCurrentStock)
There is one point to be noted though. How do we update the index?
- If we are buying with the current stock price prices[index], then we can immediately sell the next day. Hence our next call will be to index+1.
- Similarly, if we don't buy with the current stock price prices[index], then we can try buying at the next day's price. Hence our next call will be to index+1.
- The special case is when we sell at the current price. Since we cannot immediately buy and need to cool down for one day, we cannot recursively call to buy at index+1, rather we can only start buying from index+2.
- But if we don't sell at the current price, we can try to sell the next day.
But you see, the problem with this approach is that at every step, we are trying two different paths.
Hence for the stock prices with size n, the total complexity would be O(2^n) which is terrible. Can we do better? Yes, we can.

Turns out there are multiple repeated states in the problem which we can cache and use rather than re-computing.

  class Solution {
      private int maxProfitHelper(int[] prices, int index, 
                                                boolean canBuy) {
          if(index >= n) return 0;

          int price = prices[index];
          int profit = 0;

          if(canBuy) {
              int buyNow = -price + 
                           maxProfitHelper(prices, index+1, false);
              int buyNext = maxProfitHelper(prices, index+1, true);

              profit = Math.max(buyNow, buyNext);
          }
          else {
              int sellNow = price + 
                            maxProfitHelper(prices, index+2, true);
              int sellNext = 0 + 
                            maxProfitHelper(prices, index+1, false);

              profit = Math.max(sellNow, sellNext);
          }

          return memo[index][canBuy] = profit;
      }
  }

💬 Thought Process - Memoization

Since we have two states indices and canBuy, we can use a 2D array, memo[n][2], to cache values at every index and canBuy pair.
Also since we would have int values at every index, we cannot use a boolean type to indicate canBuy.
Rather, we would have to use int values to represent canBuy: 0 for false and 1 for true.
We will initialise the 2D array with -1 and for every pair of index and canBuy, we will cache the calculated values from recursive calls.
If at any call, the value at index and canBuy is not -1, which means this subproblem was already calculated we can return the solved value.

👩🏽‍💻 Solution - Memoization

class Solution {
    private int[][] memo;
    private int n;

    public int maxProfit(int[] prices) {
        n = prices.length;
        memo = new int[n][2];
        for(int[] row: memo) Arrays.fill(row, -1);

        return maxProfitHelper(prices, 0, 1);
    }

    private int maxProfitHelper(int[] prices, int index, 
                                int canBuy) {
        if(index >= n) return 0;

        if(memo[index][canBuy] != -1) return memo[index][canBuy];

        int price = prices[index];
        int profit = 0;

        if(canBuy == 1) {
            int buyNow = -price + 
                maxProfitHelper(prices, index+1, 0);
            int buyNext = 0 + 
                maxProfitHelper(prices, index+1, 1);

            profit = Math.max(buyNow, buyNext);
        }
        else {
            int sellNow = price + 
                maxProfitHelper(prices, index+2, 1);
            int sellNext = 0 +
                maxProfitHelper(prices, index+1, 0);

            profit = Math.max(sellNow, sellNext);
        }

        return memo[index][canBuy] = profit;
    }
}

Time Complexity: O(n*2)
Space Complexity: O(n*2)

💬 Thought Process - Tabulation

We can convert the above solution to tabulation format. Instead of recursive calls, we can use loops to calculate the current problem with the previously solved subproblem.
Our sup-problems would be the same. Instead of calling recursively, we can use values caches using 2D arrays.
Finally, we will have to return the dp[0][1], indicating the max profit of the first decision for buying on the stock price on day 1 or not.

👩🏽‍💻 Solution - Tabulation

class Solution {
    public int maxProfit(int[] prices) {
        int n = prices.length;
        int[][] dp = new int[n][2];

        for(int index = n-1; index >= 0; index--) {
            for(int canBuy = 0; canBuy < 2; canBuy++) {
                int price = prices[index];
                int profit = 0;

                int next = index+1;
                if(canBuy == 1) {
                    int buyNow = -price + 
                        (next >= n ? 0 : dp[next][0]);
                    int buyNext = (next >= n) ? 0 : dp[next][1];

                    profit = Math.max(buyNow, buyNext);    
                }
                else {
                    int sellNowNext = index+2;

                    int sellNow = price + 
                        (sellNowNext >= n ? 0 : dp[sellNowNext][1]);
                    int sellNext = (next >= n) ? 0 : dp[next][0];

                    profit = Math.max(sellNow, sellNext); 
                }

                dp[index][canBuy] = profit;
            }
        }

        return dp[0][1];
    }
}

Time Complexity: O(n*2)
Space Complexity: O(n*2)

You can find the link to the GitHub repo for this question here: 309. Best Time to Buy and Sell Stock with Cooldown.

Conclusion

That's a wrap for today's problem. If you liked my explanation then please do drop a like/ comment. Also, please correct me if I've made any mistakes or if you want me to improve something!

Thank you for reading!