309. Best Time to Buy and Sell Stock with Cooldown

Problem statement:

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:

prices = [1, 2, 3, 0, 2]
maxProfit = 3
transactions = [buy, sell, cooldown, buy, sell]

Solution:

This is a DP problem, I did not figure the solution, refer the idea from leetcode discussion. After thinking carefully, it makes sense. 

Also, try to summarize the DP formula by my own words.

Three arrays, buy[i], sell[i], and rest[i].

buy[i]: max profit when buying the stock at day i. It could happen in two situations: buy after cooldown or do nothing, following the value from last day.

buy[i] = max(rest[i - 1] - prices, buy[i - 1])

sell[i]: max profit when selling the stock at day i. It also could happen in two situations: sell after buy of last day or do nothing, following the value from last day.

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

rest[i]: max profit when resting at day i. It happens only in one situation: sell in last day.

rest[i] = sell[i - 1]

Do substitution in buy[i], it becomes

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

According to what described, we have to buy after one-day cooldown after a sell, translate it into DP formula: buy[i] depends on sell[i - 2](one-day cooldown).

The DP formula relies on sell[i - 2], sell[i - 1], buy[i - 1], space complexity can be reduced to O(1). Four DP variables: last_buy, buy, last_sell, sell.

last_buy and buy is initialized to INT_MIN since buy could be negative value. last_sell and sell is initialized to 0.

Only buy[i] depends on sell[i - 2], there is a tricky to represent sell[i - 2]. It is about when to update last_buy, buy, last_sell, sell.

Each time enters a loop, sell and buy become last_sell and last_buy(sell[i - 1], buy[i - 1]), last_sell and last_buy become sell[i - 2] and buy[i - 2]. 

• First, update last_buy(becomes buy[i - 1]) 
• Second, update buy(becomes buy[i]) using last_buy and last_sell(it is sell[i - 2] now).  
• Third, update last_sell(becomes sell[i - 1]) and sell(becomes sell[i]).

Return sell. 

Time complexity O(n).

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int last_sell = 0;
        int sell = 0;
        int last_buy = INT_MIN;
        int buy = INT_MIN;
        for(auto price : prices){
            last_buy = buy;
            buy = max(last_sell - price, last_buy);
            last_sell = sell;
            sell = max(last_buy + price, last_sell);
        }
        return sell;
    }
};