Best Time to Buy and Sell Stock with Cooldown

1. Problem Description

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

Design an algorithm to compute the maximum profit. Under the following constraints, you may complete as many transactions as possible (buying and selling the same stock multiple times):

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

2. Approach

1. Approach 1
   • DFS
   • Assume the stock is bought on day i, then find a later day j with price > price[i] and calculate the profit
   • Continue from j+2

• DFS + memoization

3. Implementation

3.1. DFS + Memoization (Timeout)

func maxProfit(prices []int) int {
    memo := make(map[int]int, 0)
    return maxProfitDFS(prices, 0, memo)
}

func maxProfitDFS(prices []int, index int, memo map[int]int) int {
    profit, ok := memo[index]
    if ok {
        return profit
    }

    if index >= len(prices) {
        return 0
    }

    profit = 0
    for i := index; i < len(prices)-1; i++ {
        buy := prices[i]
        for j := i+1; j < len(prices); j++ {
            sell := prices[j]
            if sell > buy {
                profit = max(profit, sell-buy+maxProfitDFS(prices, j+2, memo))
            }
        }
    }
    memo[index] = profit
    return profit
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

3.2. DFS + Memoization

const (
    OptBuy = 0
    OptSell = 1
)

func maxProfit(prices []int) int {
	memo := make(map[string]int, 0)
	return maxProfitDFS(prices, 0, OptBuy, memo)
}

// Profit obtained by performing opt on day index
func maxProfitDFS(prices []int, index int, opt int, memo map[string]int) int {
	if index >= len(prices) {
		return 0
	}

	if profit, ok := memo[getKey(index, opt)]; ok {
		return profit
	}

	// Three variables record [do nothing], [buy], and [sell]
	a, b, c := 0, 0, 0
	// Do nothing
	a = maxProfitDFS(prices, index+1, opt, memo)
	// Sell
	if opt == OptSell {
		// After selling there is a one-day cooldown, so the next operation can only happen at index+2
		b = maxProfitDFS(prices, index+2, OptBuy, memo) + prices[index]
	// Buy
	} else {
		// There is no cooldown after buying, so the next operation can happen at index+1
		c = maxProfitDFS(prices, index+1, OptSell, memo) - prices[index]
	}

	// The final result is the maximum of the three variables
	profit := Max(a, b, c)
	memo[getKey(index, opt)] = profit
	return profit
}

func getKey(index int, opt int) string{
	return fmt.Sprintf("%v:%v", index, opt)
}

func Max(data ...int) int {
	max := data[0]
	for i := 1; i < len(data); i++ {
		if data[i] > max {
			max = data[i]
		}
	}
	return max
}

4. References

• 309. Best Time to Buy and Sell Stock with Cooldown - LeetCode