309. Best Time to Buy and Sell Stock with Cooldown

Problem

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]

Problem Summary

Given an integer array, where the i-th element represents the stock price on day i.​

Design an algorithm to calculate the maximum profit. Under the following constraints, you may complete as many transactions as possible (buy and sell one stock multiple times):

• You cannot participate in multiple transactions at the same time (you must sell the stock before buying again).
• After selling the stock, you cannot buy stock on the next day (i.e., the cooldown period is 1 day).

Solution Approach

• Given an array representing the price of a stock on each day. Design a trading algorithm to automatically trade over these days, with the requirements that: only one operation can be performed each day; after buying a stock, you must sell it before buying again; after each sale, you cannot buy on the next day. How should you trade to maximize profit?
• This problem is a variation of Problems 121 and 122.
• There are 3 possible operations each day: buy, sell, and cooldown. The day after a sell must be cooldown, but cooldown can occur on any day. For example: buy，cooldown，cooldown，sell，cooldown，cooldown. buy[i] represents the maximum profit obtainable by ending day i with a buy or cooldown. For example: buy, sell, buy or buy, cooldown, cooldown. sell[i] represents the maximum profit obtainable by ending day i with a sell or cooldown. For example: buy, sell, buy, sell or buy, sell, cooldown, cooldown. price[i-1] represents the stock price on day i (because price starts from 0).
• If day i is sell, then the maximum profit obtainable on this day is buy[i - 1] + price[i - 1], because you can sell only after buying. If this day is cooldown, then the maximum profit obtainable on this day is still sell[i - 1]. Therefore, the state transition equation for sell[i] is sell[i] = max(buy[i - 1] + price[i - 1], sell[i - 1]). sell[0] = 0 means selling on the first day; since no stock is held on the first day, sell[0] = 0. sell[1] = max(sell[0], buy[0]+prices[1]) means comparing selling on the first day with not selling on the first day and selling on the second day, and storing the larger amount in sell[1].
• If day i is buy, then the maximum profit obtainable on this day is sell[i - 2] - price[i - 1], because day i - 1 is cooldown. If this day is cooldown, then the maximum profit obtainable on this day is still buy[i - 1]. Therefore, the state transition equation for buy[i] is buy[i] = max(sell[i - 2] - price[i - 1], buy[i - 1]). buy[0] = -prices[0] means buying on the first day, so the money becomes negative. buy[1] = max(buy[0], -prices[1]) means not buying on the first day and buying on the second day instead.

Code

package leetcode

import (
	"math"
)

// Solution 1 DP
func maxProfit309(prices []int) int {
	if len(prices) <= 1 {
		return 0
	}
	buy, sell := make([]int, len(prices)), make([]int, len(prices))
	for i := range buy {
		buy[i] = math.MinInt64
	}
	buy[0] = -prices[0]
	buy[1] = max(buy[0], -prices[1])
	sell[1] = max(sell[0], buy[0]+prices[1])
	for i := 2; i < len(prices); i++ {
		sell[i] = max(sell[i-1], buy[i-1]+prices[i])
		buy[i] = max(buy[i-1], sell[i-2]-prices[i])
	}
	return sell[len(sell)-1]
}

// Solution 2 DP with optimized auxiliary space
func maxProfit309_1(prices []int) int {
	if len(prices) <= 1 {
		return 0
	}
	buy := []int{-prices[0], max(-prices[0], -prices[1]), math.MinInt64}
	sell := []int{0, max(0, -prices[0]+prices[1]), 0}

	for i := 2; i < len(prices); i++ {
		sell[i%3] = max(sell[(i-1)%3], buy[(i-1)%3]+prices[i])
		buy[i%3] = max(buy[(i-1)%3], sell[(i-2)%3]-prices[i])
	}
	return sell[(len(prices)-1)%3]
}