Update 06. With Cooldown.py

02_Dynamic-Programming/11. Buy and Sell Stock Problems/06. With Cooldown.py

@@ -1,23 +1,18 @@
-# https://youtu.be/GY0O57llkKQ
 # https://leetcode.com/problems/best-time-to-buy-and-sell-stock-with-cooldown/
 
 class Solution:
-    def maxProfit(self, prices):
-        obsp = - prices[0]  # Old Brought State Profit
-        ossp = 0            # Old Sold State Profit
-        ocsp = 0            # Old Cold State Profit
-
-        for i in range(1, len(prices)):
-            nbsp = 0        # New Brought State Profit
-            nssp = 0        # New Sold State Profit
-            ncsp = 0        # New Cold State Profit
-
-            nbsp = max(ocsp - prices[i], obsp)
-            nssp = max(obsp + prices[i], ossp)
-            ncsp = max(ossp, ocsp)
-
-            obsp = nbsp
-            ossp = nssp
-            ocsp = ncsp
-
-        return ossp
+    def maxProfit(self, prices: List[int]) -> int:
+        memo = {}
+        def solve(i, can_sell):
+            if i >= len(prices): return 0
+            if (i,can_sell) in memo: return memo[(i,can_sell)]
+            if can_sell == 1: 
+                profit = max(prices[i] + solve(i+2,0), solve(i+1,1))    # After selling cooldown one day
+            else:
+                profit = max(-prices[i] + solve(i+1,1), solve(i+1,0))
+            memo[(i,can_sell)] = profit
+            return profit
+        return solve(0,0)
+
+# Time: O(2 * N)
+# Space: O(2 * N)