🌱 Keegan's Digital Garden

309. Best Time to Buy and Sell Stock with Cooldown

4 min read

📝 Problem Details

🧪Problem Description

You are given an integer array `prices` where `prices[i]` is the price of a given stock on the `i`th day.

Find the maximum profit you can achieve. You may complete as many transactions as you like 
(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 buying again).

🔭Key Observations

1. The problem requires maximizing the profit with two main states:
   - Buying state: Whether you are looking to buy a stock on a specific day.
   - Selling state: Whether you are looking to sell a stock on a specific day.
   
2. Cooldown Mechanism:
   - If you sell on day `i`, you cannot buy again until day `i+2`.
   
3. Recursive Substructure:
   - If in the buying state at day `i`, you can either:
     a) Buy the stock on day `i` and move to the selling state for day `i+1`.
     b) Skip buying (cooldown) and remain in the buying state for day `i+1`.
   - If in the selling state at day `i`, you can either:
     a) Sell the stock on day `i` and move to the buying state for day `i+2` (due to cooldown).
     b) Skip selling (cooldown) and remain in the selling state for day `i+1`.

4. Optimal Substructure:
   - The problem can be solved using recursion with memoization as overlapping subproblems exist.
   
5. Base Case:
   - When `i >= prices.size()`, there is no profit to gain, so return `0`.

💡 Explanation of Solution

1. **Dynamic Programming Approach:**
   - Define a recursive function `dfs(i, buying)` where:
     - `i` is the current day index.
     - `buying` is a boolean indicating whether you are in the buying state.
   - Use memoization (hashmap `dp`) to store the maximum profit for each state `(i, buying)`.

2. **State Transitions:**
   - If in the buying state on day `i`:
     a) **Buy:** Transition to day `i+1` in the selling state, reducing profit by `prices[i]`.
     b) **Cooldown:** Transition to day `i+1` in the buying state without buying.
     - Take the maximum of these two options.
   - If in the selling state on day `i`:
     a) **Sell:** Transition to day `i+2` in the buying state, increasing profit by `prices[i]`.
     b) **Cooldown:** Transition to day `i+1` in the selling state without selling.
     - Take the maximum of these two options.

3. **Implementation Details:**
   - Use a `unordered_map` with a custom hash function to memoize results for pairs `(i, buying)`.
   - Use recursion to compute results bottom-up, starting from day `0` in the buying state.

4. **Return the Result:**
   - Call the `dfs(0, true)` function to calculate the maximum profit starting on day 0 in the buying state.

⌛ Complexity Analysis

1. **Time Complexity:**
   - There are `O(n)` unique states for each combination of `i` (days) and `buying` (true/false).
   - Each state is computed once due to memoization.
   - Overall complexity: `O(n)`.

2. **Space Complexity:**
   - Space is used for the recursion stack and memoization:
     a) Recursion stack: `O(n)` in the worst case.
     b) Memoization table: `O(n)` states.
   - Overall space complexity: `O(n)`.

💻 Implementation of Solution

class Solution {
public:
    int maxProfit(vector<int>& prices) {
        // State: buying or selling?
        // Buying --> i+1
        // Selling --> i+2 (mandatory cooldown day)
 
        // Hash function for pair<int, bool>
        auto pair_hash = [](const pair<int, bool>& p) -> size_t {
            return hash<int>()(p.first) ^ (hash<bool>()(p.second) << 1);
        };
 
        // Memoization table with custom hash function
        unordered_map<pair<int, bool>, int, decltype(pair_hash)> dp(0, pair_hash);
 
        // Recursive function to compute max profit
        function<int(int, bool)> dfs = [&](int i, bool buying) -> int {
            // Base case: if index is out of bounds
            if (i >= prices.size()) return 0;
 
            // Check if state is already computed
            if (dp.find({i, buying}) != dp.end()) return dp[{i, buying}];
 
            // Compute result based on current state
            if (buying) {
                int buy = dfs(i + 1, false) - prices[i]; // Buy stock
                int cooldown = dfs(i + 1, true);        // Skip (cooldown)
                dp[{i, buying}] = max(buy, cooldown);
            } else {
                int sell = dfs(i + 2, true) + prices[i]; // Sell stock
                int cooldown = dfs(i + 1, false);       // Skip (cooldown)
                dp[{i, buying}] = max(sell, cooldown);
            }
 
            return dp[{i, buying}];
        };
 
        // Start recursion from day 0 in the buying state
        return dfs(0, true);
    }
};

Graph View