Best time to buy and sell stock with cooldown

Shreya Deep
Last Updated: May 13, 2022

Introduction

Stocks are a great way to make profits. If you buy a stock at price x and sell it at a price y, then the profit is (y-x). That’s why to make a positive profit, we wish to sell the stock at a price higher than the price at which we bought it. Otherwise, it would be a loss. One whole buy + sell action is called a transaction. 

In this article, we will discuss the problem of finding the maximum profit you can achieve by buying and selling stocks based on the given constraints.

Problem Statement

Given the stock prices of n days in an array “prices,” where prices[i] denotes the price of the stock at day i. Find the maximum profit you can achieve by doing as many transactions as you like provided that: 

After you sell the stock, you cannot buy the stock the next day, i.e., you have a cooldown period of 1 day.

Note: You’re not allowed to be engaged in more than one transaction at the same time. It means you must sell the stock before buying again.

For example:

INPUT : prices = [1,2,3,0,2]

OUTPUT:  3

Buy on day one and sell on the second day. Profit = 2-1 = 1.

The 3rd day will be a cooldown day.

Again buy on day four and sell on the fifth day. Profit = 2-0 = 2.

Total profit = 1+2 = 3.

INPUT : prices = [1]

OUTPUT:  0

For one complete transaction, we need at least two days. Here is just one day, we can’t make any transaction, so the profit is 0.

Approach 1: Recursion

In this problem, before starting a new transaction, we’re first required to sell the last stock we bought. Now let’s assume that we buy a stock on day i. To sell this stock, we need to find a day j, such that j>i and j<n. Now let’s assume that there are many such j’s. So the question is which one should we choose? Now see that if we decide on one such j, then for the subsequent transactions, we have reduced our search space to j+1 to n-1. 

Again, we need to do the same process to our new search space. Therefore, recursion can be used in this problem. 

The idea is that on any day, we have two options regarding buying the stock. Either we buy it on that day, or we don’t. We maintain a bool variable “canBuy,” which tells us this. If we can buy stock on day i, canBuy will be 1 otherwise, it will be 0. 

Now, if we can buy on day i, we further have two options :

  • Simply buy the stock on day i so the profit will be reduced by the cost of stock at that index and then look for the profit from day i+1 with canBuy = 0.
  • Or, don’t buy the stock on day i, rather, move to day i+1 with canBuy=1. 

If we have to sell the stock on day i, we further have two options:

  • Sell the stock on day i and add the cost of the stock at this day to the profit. Now, since we’ve sold the stock on day i, i+1 will be a cooldown day and therefore, next what we can do is, move to day i+2 with canBuy = 1.
  • Or, don’t sell the stock on day i, rather, move to day i+1 with canBuy = 0.

We’ll make a function named “func” which will do the above calculations for day i.

Steps are: 

  • Input the number of days, say n, and the prices on all n days.
  • Call the maxProfit function with the parameter prices, which will basically return the final answer. 
  • In the maxProfit function, since on the first day we can only buy the stock, Call the “func” function with canBuy ==true for day 0. 
  • In the “func” function, check if we’ve reached the end of proces vector. If yes, return 0 as no further profit can be made.
  • Otherwise, if canBuy is true, which means we can buy stock on day i, calculate the profit for the two subproblems and return the maximum of these two. The subproblems are - 
    • Simply buy the stock on day i so the profit will be reduced by the cost of stock at that index and then call “func” for index i+1 with canBuy = 0, i.e., call func(prices, i+1,0) and subtract prices[i] from it.
    • Or, don’t buy the stock on day i, rather, move to day i+1 with canBuy=1 i.e., call func(prices, i+1,1).
    • Return max of the above two.
  • If the canBuy is false, again explore the two options for false condition discussed above and return the maximum of those two. The subproblems are - 
    • Sell the stock on day i and add the cost of the stock at this day to the profit. Call func for i+2 with canBuy=1, as i+1th day will be the cooldown day, i.e., call func(prices,i+2,1) and add prices[i] to it. 
    • Or, don’t sell the stock on day i, rather, move to day i+1 with canBuy = 0 i.e., call func(prices,i+1,0).
    • Return max of the above two.
  • maxProfit will return the answer to the main function. So just print it. 

C++ implementation

Below is the C++ implementation of the above discussed approach.

#include <bits/stdc++.h>
using namespace std;

int func(vector<int>& prices,int i,bool canBuy){
    if(i>=prices.size())  // If we reach the end of the vector, return 0 as no further profit can be made.
        return 0;
    int x,y;
    if(canBuy==1){ // if we can buy, we have these two options. Return the maximum of these two
        x=func(prices,i+1,0)-prices[i];  
        y=func(prices,i+1,1);
        return max(x,y);
    }
    else{// if we can sell, we have these two options. Return the maximum of these two.
        x=prices[i]+func(prices,i+2,1);
        y=func(prices,i+1,0);
        return max(x,y);
    }
}

int maxProfit(vector<int>& prices) {
    int idx=0;
    return func(prices,idx,1); // Starting with index 0 call func with canBuy = 1.
}

int main(){
    int n;
    cin>>n; // number of days
    vector<int>prices(n);  // Declaration of prices vector
    for(int i=0;i<n;i++){  //inputting the values of the prices vector
        cin>>prices[i]; 
    }
    int ans = maxProfit(prices); 
    cout<<ans<<endl;
    return 0;
}

Complexities

Time complexity

O(2^n) , where n is the number of days.

Reason: In the worst case, for every index, we make two recursive calls as we have two options. The maximum depth of the recursion tree can go up to n. Hence the time complexity is O(2 ^ n).

Space complexity

O(n), where n is the number of days.

Reason: In the worst case, the extra space used by the recursion stack can go up to a  maximum depth of n. Hence the space complexity is O(n).

Approach 2: Dynamic Programming

The time complexity of the above approach is exponential and thus will not pass all test cases. In the above recursive approach, we have overlapping subproblems and optimal substructure. Therefore, we can optimize it using dynamic programming.

What we do is memorize the results of each pair of index i and canBuy state in a dp lookup table. The dp table will be a 2-D table with size [n]x[2] as there can be n indexes and the canBuy variable has maximum two different types of values (0 or 1). 

If we call the function “func” with some state (i,canBuy) which has already been calculated, we just simply return the memorized solution. This will help reduce the time complexity as we don’t do repeated calculations anytime. For this, initially in the dp table we initialize the values with -1 and we keep updating the answer for current (i,canBuy) state. If there is a state whose dp value is not equal to -1, this means that we’ve already calculated the value for this state. So we just return the memorized dp value. 

Rest of the whole idea is the same.

Steps are: 

  • Input the number of days, say n, and the prices on all n days.
  • Call the maxProfit function with the parameter prices, which will basically return the final answer. 
  • In the maxProfit function, since on the first day we can only buy the stock, Call the “func” function with canBuy ==true for day 0. 
  • In the “func” function, check if we’ve reached the end of proces vector. If yes, return 0 as no further profit can be made.
  • Check if the current (i,canBuy) state already has an answer calculated, i.e., see if dp[i][canBuy] is != -1. If true, then the answer is already calculated, so just return it. Otherwise, continue with the following steps.
  • if canBuy is true, which means we can buy stock on day i, calculate the profit for the two subproblems and return the maximum of these two. The subproblems are - 
    • Simply buy the stock on day i so the profit will be reduced by the cost of stock at that index and then call “func” for index i+1 with canBuy = 0, i.e., call func(prices, i+1,0) and subtract prices[i] from it.
    • Or, don’t buy the stock on day i, rather, move to day i+1 with canBuy=1 i.e., call func(prices, i+1,1).
    • Return and store the max of the above two in dp[i][canBuy].
  • If the canBuy is false, again explore the two options for false condition discussed above and return the maximum of those two. The subproblems are - 
    • Sell the stock on day i and add the cost of the stock at this day to the profit. Call func for i+2 with canBuy=1, as i+1th day will be the cooldown day, i.e., call func(prices,i+2,1) and add prices[i] to it. 
    • Or, don’t sell the stock on day i, rather, move to day i+1 with canBuy = 0 i.e., call func(prices,i+1,0).
    • Return and store the max of the above two in dp[i][canBuy].
  • maxProfit will return the answer to the main function. So just print it. 

 C++ implementation

Below is the C++ implementation of the above discussed approach.

#include <bits/stdc++.h>
using namespace std;

int func(vector<int>& prices,int i,bool canBuy,vector<vector<int>>&dp){
    if(i>=prices.size())  // If we reach the end of the vector, return 0 as no further profit can be made.
        return 0;
    if(dp[i][canBuy]!=-1){ //Value for this state is already calculated, so just return it.
        return dp[i][canBuy] ;
    }
    int x,y;
    if(canBuy==1){ // if we can buy, we have these two options. Return the maximum of these two
        x=func(prices,i+1,0,dp)-prices[i];  
        y=func(prices,i+1,1,dp);
        return dp[i][canBuy] = max(x,y);
    }
    else{// if we can sell, we have these two options. Return the maximum of these two.
        x=prices[i]+func(prices,i+2,1,dp);
        y=func(prices,i+1,0,dp);
        return dp[i][canBuy] = max(x,y);
    }
}

int maxProfit(vector<int>& prices) {
    int idx=0;
    int n = prices.size();
    vector<vector<int>>dp(n,vector<int>(2,-1));
    return func(prices,idx,1,dp); // Starting with index 0 call func with canBuy = 1.
}

int main(){
    int n;
    cin>>n; // number of days
    vector<int>prices(n);  // Declaration of prices vector
    for(int i=0;i<n;i++){  //inputting the values of the prices vector
        cin>>prices[i]; 
    }
    int ans = maxProfit(prices); 
    cout<<ans<<endl;
    return 0;
}


Complexities

Time complexity

O(n) , where n is the number of days.

Reason: For every i, we’re calculating the value of dp[i][canBuy] only once. We are calling the function func for a repeated number of times but do the calculation only once.

Space complexity

Reason: O(2*n), where n is the number of days.

The extra is the space used by the dp table. Since the size of the dp table is 2*n, the space complexity is O(2*n).

Approach 3: Constant space solution

The space complexity of the previous solution was O(n). So, if we can optimize it to space complexity of O(1), we should surely do it. Let’s see how we can do that.

Let’s define: 

buy[i]: the maximum profit up to day i with buy as last action;

sell[i]: the maximum profit up to day i with sell as last action;

In that case, the recurrence relations will be:

  1. buy[i] = max(sell[i-2]-price[i], buy[i-1])  

Reason: buy[i] will be the profit with buy as last action. So on day i, if we don’t do any of the actions, i.e, if we take rest, it means we must have done the “buy” action on some previous day. So buy[i-1] can give the information about that. In case we actually do the “buy” action on the day i, the profit will be sell[i-2]-price[i] because i-1th  day will be the cooldown. 

  1. sell[i] = max(buy[i-1]+price[i], sell[i-1])

Reason: sell[i] will be the profit with sell as last action. So on day i, if we don’t do any of the actions, i.e, if we take rest, it means we must have done the “sell” action on some previous day. So sell[i-1] can give the information about that. In case we actually do the “sell” action on the day i, the profit will be buy[i-1]+price[i] because we must have brought the stock on some day before i so that we’re able to sell it on day i. 

Thus we can see that with “buy” as the last action, the profit on day i depends on the profit of day i-1 and selling profit of day i-2. Similarly, with “sell” as the last action, the profit on day i depends on the profit of day i-1 and buying profit of day i-1. So, if we store these 5 values in some variables, our work is done!!

Steps are: 

  • We initialize 5 variables. 
    • buy which denotes buy[i] to 0
    • buyPre which denotes buy[i-1] to INT_MIN 
    • sell which denotes sell[i] to 0
    • sellPre1 which denotes sell[i-1] to 0
    • sellPre2 which denotes sell[i-2] to 0
  • Traverse the prices vector. For each price p in the vector,
    • Update buy = max(buyPre,sellPre2-p)
    • Update sell = max(sellPre1,buyPre+p)
    • For the next iteration, update the values of buyPre=buy, sellPre2 = sellPre1, sellPre1 = sell.
  • Return the value of sell as the last action we must have done would be sell.

C++ implementation

Below is the C++ implementation of the above-discussed approach.

#include <bits/stdc++.h>
using namespace std;

int maxProfit(vector<int>& prices) {
    int buy /*buy[i]*/, buyPre = INT_MIN /*buy[i-1]*/;
      int sell = 0 /*sell[i]*/, sellPre1 = 0 /*sell[i-1]*/, sellPre2 = 0 /*sell[i-2]*/;
      
      for (int p : prices) {
        // updating the values of buy and sell for the current p
        buy = max(buyPre, sellPre2 - p);
        sell = max(sellPre1, buyPre + p);
        
       // update for next iteration
        buyPre = buy;
        sellPre2 = sellPre1;
        sellPre1 = sell;
      }
      
      return sell;
}

int main(){
    int n;
    cin>>n; // number of days
    vector<int>prices(n);  // Declaration of prices vector
    for(int i=0;i<n;i++){  //inputting the values of the prices vector
        cin>>prices[i]; 
    }
    int ans = maxProfit(prices); 
    cout<<ans<<endl;
    return 0;
}


Complexities

Time complexity

O(n), where n is the number of days.

Reason: We’re just iterating through the whole prices vector once.

Space complexity

O(1)

Reason: The only spaces taken here are by the variables that just take constant space. 

Frequently asked questions

  1. What is dynamic programming, and where is it used?
    Dynamic programming is an optimization method used in various programming problems. It is used in problems where the solution depends on smaller overlapping subproblems. We use it to memorize the results so that they can easily be used later when needed.
     
  2. What are overlapping subproblems?
    A problem has overlapping subproblems if it can be divided into smaller problems that are reused multiple times.
     
  3. What is optimal substructure?
    A problem is said to have optimal substructure if an optimal solution can be constructed from optimal solutions of its subproblems.
     
  4. When does a stock transaction give profit?
    When you sell the stock at a price that is higher than the price at which you bought it, the stock transaction gives profit.
     
  5. Are there more Data Structures and Algorithms problems in CodeStudio?
    Yes, CodeStudio is a platform that provides both practice coding questions and commonly asked interview questions. The more we’ll practice, the better our chances are of getting into a dream company of ours.

Key Takeaways

In this article, we’ve discussed the best time to buy and sell stock with cooldown problem. There are various problems similar to this such as Best time to buy and sellBuy and sell stockSelling stockBest time to buy and sell stock, and Buy and sell stock - III. I would suggest you solve them to gain more confidence in these kinds of problems. These questions are asked during various coding contests as well as placements tests and thus are very important. 

To practice more such problems, Codestudio is a one-stop destination. This platform will help you acquire effective coding techniques and give you an overview of student interview experience in various product-based companies.

Happy Coding!

Was this article helpful ?
0 upvotes