< Back





309. Best Time to Buy and Sell Stock with Cooldown

Fun problem, two great ways to solve it. My first solution wasn't space optimal, however, it's easy
to reason about recursively. Like every dynamic programming problem, we have states and we have to
make decisions at each state, and we use the previous solutions / decisions to inform our current
decision so we don't have to recalculate values.

In our first solution, we recursively choose 4 options:

  * Skip, making no choice
  * If we hold, we continue to hold or sell and transition to the cooldown state
  * If we're in the cooldown state, we stay in the cooldown state or transition to the ready state
  * If we're free, we remain free or transition to the hold state

Using the @cache decorator for our dynamic programming function, we use memoization to avoid
recalculating already visted paths in this search. Eventually, we find the maximum.

Alternatively, we represent this state machine with a different algorithm, using the same states,
and we update the best choices for each possibility after inspecting each price.

The recursive solution is as follows:


  class Solution:
      def maxProfit(self, prices: List[int]) -> int:
          n = len(prices)

          @cache
          def dp(i: int, h: bool, c: bool) -> int:
              if i == n:
                  return 0

              ans = dp(i + 1, h, c)

              if h:
                  ans = max(ans, prices[i] + dp(i + 1, False, True))
              elif c:
                  ans = max(ans, dp(i + 1, h, False))
              else:
                  ans = max(ans, -prices[i] + dp(i + 1, True, False))

              return ans

          return dp(0, False, False)


_ Time Complexity:

  O(n) - Where n is the number of prices.

_ Space Complexity:

  O(n) - We create a recursive call stack n times.


The iterative solution is as follows:


  class Solution:
      def maxProfit(self, prices: List[int]) -> int:
          sold, held, cooldown = float("-inf"), float("-inf"), 0

          for price in prices:
              sold, held, cooldown = (
                  held + price,
                  max(held, cooldown - price),
                  max(cooldown, sold),
              )

          return max(sold, cooldown)


_ Time Complexity:

  O(n) - Where n is the number of prices.

_ Space Complexity:

  O(1) - We store our states in constant space.