Best Time to Buy and Sell Stock with Cooldown

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]

Solutions

Let stock[i + 1] be the maximum profit at day i holding stock. Let money[i + 1] be the maximum profit at day i without holding stock. (days start from 0).

To have stock at day i, we can either:

• have stock at day i-1, then the profit is stock[i]; or

• buy stock at day i, then we must not sell at day i-1. The profit is money[i-1] - prices[i].

• Thus, stock[i + 1] = max(stock[i], money[i - 1] - prices[i]).

To not have stock at day i, we can either:

• don't have stock at day i-1 and don't buy at day i, then the profit is money[i-1]; or

• have stock at day i-1 and sell the stock at day i. Then the profit is stock[i-1] + prices[i].

• Thus, money[i + 1] = max(stock[i] + prices[i], money[i]).

money[i] always larger than stock[i], so we return money[n].

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int n = prices.size();
        if(n == 0) return 0;
        vector<int> stock(n + 1);
        vector<int> money(n + 1);
        stock[1] = -prices[0];
        money[1] = 0;
        for(int i = 1; i < n; ++i){
            stock[i + 1] = max(stock[i], money[i - 1] - prices[i]);
            money[i + 1] = max(stock[i] + prices[i], money[i]);
        }
        return money[n];
    }
};

Optimization

We can optimize to algorithm to use constant space.

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        int n = prices.size();
        if(n == 0) return 0;
        int stock = -prices[0];
        int money = 0;
        int prev_money = 0;
        for(int i = 1; i < n; ++i){
            // before:
            // stock = stock[i]
            // prev_money = money[i - 1]
            // after:
            // stock = stock[i + 1]
            stock = max(stock, prev_money - prices[i]);
            // money = money[i]
            prev_money = money;
            // before:
            // money = money[i]
            // stock = stock[i + 1]
            // if stock[i + 1] = stock[i]
            // <==> stock[i] > money[i - 1] - prices[i]
            // then it does the same as the previous solution
            // else stock[i] < money[i - 1] - prices[i]
            // ==> stock[i + 1] = money[i - 1] - prices[i]
            // ==> stock + prices[i] = stock[i + 1] + prices[i]
            // = money[i - 1] - prices[i] + prices[i + 1]
            // = money[i - 1] <= money = money[i]
            // thus money[i + 1] must be money[i]
            money = max(stock + prices[i], money);
        }
        return money;
    }
};