Dynamic Programming - Best Time to Buy and Sell Stock with Cooldown
All diagrams presented herein are original creations, meticulously designed to enhance comprehension and recall. Crafting these aids required considerable effort, and I kindly request attribution if this content is reused elsewhere.
Difficulty : Easy
DFS
Problem
You are given an array prices where prices[i] is the price of a given stock on the ith day.
Find the maximum profit you can achieve. You may complete as many transactions as you like (i.e., buy one and sell one share of the stock multiple times) with the following restrictions:
- After you sell your stock, you cannot buy stock on the next day (i.e., cooldown one day).
Note: You may not engage in multiple transactions simultaneously (i.e., you must sell the stock before you buy again).
Example 1:
1
2
3
Input: prices = [1,2,3,0,2]
Output: 3
Explanation: transactions = [buy, sell, cooldown, buy, sell]
Example 2:
1
2
Input: prices = [1]
Output: 0
Solution
We will be solving this problem using DFS. There are 2 options in each turn.
- If can Buy then Buy or cool down
- If can Sell then Sell or cool down.
We need to traverse in both the directions, either buy/cool down or sell/cool down and selected which ever provides max profit. Since we are going to use dfs() the max will be starting from leaf nodes at the bottom of the tree and propagate upward.
Here is the visualization for the Example 1. This shows only the partial tree, the max gets propagated from the leaf nodes to the root node.
Lets start by defining our dfs() function. Since the path we can take depends on if we can_buy or not (that is not can_buy means sell) and this is going to be different in each node, let’s pass that to the dfs() function along with the index parameter.
1
def dfs(i, can_buy):
At anytime, if the index i reached the end of the prices then we know we cannot make any more profit, so we return 0.
1
2
3
def dfs(i, can_buy):
if i >= len(prices):
return 0
Now if we are eligible to buy then we can either buy or cool down. Call dfs() by incrementing the index and setting can_buy to False. We also need to subtract the price from returned profit.
1
2
if can_buy:
profit_if_buying_now = dfs(i+1, False) - prices[i]
We also need to traverse the path to cool down instead of buying.
1
profit_if_cooldown =dfs(i+1, True)
Now return the max of both.
1
return max(profit_if_buying_now,profit_if_cooldown)
Alternatively, if we can sell then sell and skip the next element as we need to cool down after selling. Also add the price to profit.
1
2
3
4
else:
profit_if_selling = dfs (i+2, True) + prices[i]
profit_if_cooldown =dfs(i+1, False)
return max(profit_if_selling,profit_if_cooldown)
Now call the dfs() function.
1
return dfs(0, True)
The above code will work fine, however will get timeout error without memorization/caching. Let’s implement that. The dfs() function takes two arguments, so that simplest approach is to create a map using both arguments as keys.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
cache = {}
def dfs(i, can_buy):
if i >= len(prices):
return 0
if (i, can_buy) in cache:
return cache[(i, can_buy)]
if can_buy:
profit_if_buying_now = dfs(i+1, False) - prices[i]
profit_if_cooldown =dfs(i+1, True)
cache[(i, can_buy)] = max(profit_if_buying_now,profit_if_cooldown)
else:
profit_if_selling = dfs (i+2, True) + prices[i]
profit_if_cooldown =dfs(i+1, False)
cache[(i, can_buy)] = max(profit_if_selling,profit_if_cooldown)
return cache[(i, can_buy)]
Final Code
Here is the full code.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def maxProfit(prices: List[int]) -> int:
cache = {}
def dfs(i, can_buy):
if i >= len(prices):
return 0
if (i, can_buy) in cache:
return cache[(i, can_buy)]
if can_buy:
profit_if_buying_now = dfs(i+1, False) - prices[i]
profit_if_cooldown =dfs(i+1, True)
cache[(i, can_buy)] = max(profit_if_buying_now,profit_if_cooldown)
else:
profit_if_selling = dfs (i+2, True) + prices[i]
profit_if_cooldown =dfs(i+1, False)
cache[(i, can_buy)] = max(profit_if_selling,profit_if_cooldown)
return cache[(i, can_buy)]
return dfs(0, True)