Best Time to Buy and Sell Stock with Cooldown

Patrick Leaton
Published: June 3, 2021, 3:46 a.m.
Problem Description

Say you have an array for which the ith element is the price of a given stock on day i.

Design an algorithm to find the maximum profit. You may complete as many transactions as you like (ie, buy one and sell one share of the stock multiple times) with the following restrictions:

  • You may not engage in multiple transactions at the same time (ie, you must sell the stock before you buy again).
  • After you sell your stock, you cannot buy stock on next day. (ie, cooldown 1 day)

Example:

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

 

The description was taken from https://leetcode.com/problems/best-time-to-buy-and-sell-stock-with-cooldown/.

Problem Solution

#O(N) Time, O(1) Space
class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        sold, held, refresh = float('-inf'), float('-inf'), 0
        for price in prices:
            cooling = sold 
            sold = held + price
            held = max(held, refresh - price)
            refresh = max(refresh, cooling)
        return max(sold, refresh)

Problem Explanation

We can solve this problem by calculating a few different scenarios in parallel during each day and seeing which combination of scenarios gives us the best bank at the end.

We will start by setting sold, the price we were holding stock at plus the given day's price, held, the price we were holding a given stock at, and refresh, which is for undoing a previous action.

We will want to start sold and held off at negative infinity because there's no way we could show up on day one already holding or selling a stock that was bought on a previous day.

        sold, held, refreshed = float('-inf'), float('-inf'), 0

 

We will traverse each day's stock prices.

        for price in prices:

 

Each day we will want to keep track of a few different scenarios, also known states.

Cooling means that we sold yesterday and are using today as a cooldown period.

            cooling = sold 

 

Sold means that we were holding onto a stock from a previous day and we will sell it at the current day's price.

            sold = held + price

 

Held means that we are either going to keep holding a previous stock, or we refreshed yesterday and we are going to buy at the current stock price today depending on the greater value between the two choices.

            held = max(held, refresh - price)

 

Refresh means we are either going to keep the previously refreshed price or we showed up to work hungover yesterday and sold at too low of a price so the cooling value is low, and we would like to refresh to undo our last sell, depending on the greater value between the two choices.

            refresh = max(refresh, cooling)

 

Once we have finished a week's worth of work at the stock market, either the last deal we made was a good sell or it sucked so we would want to keep the last refresh price.

        return max(sold, refresh)