feature: best time to buy and sell stock with cooldown top down revisit

dynamic_programming/BestTimeToBuyAndSellStockWithCooldown.py

@@ -1,15 +1,19 @@
-from functools import cache
+from functools import lru_cache
 from typing import List
 
 
 # O(n) time || O(n) space
-def max_profit(self, prices: List[int]) -> int:
-    @cache
-    def dp(i, buy):
+def max_profit_top_down(self, prices: List[int]) -> int:
+    @lru_cache(None)
+    def dp(i, holding):
         if i >= len(prices):
             return 0
 
-        cooldown = dp(i + 1, buy)
-        return max(dp(i + 1, not buy) - prices[i], cooldown) if buy else max(dp(i + 2, not buy) + prices[i], cooldown)
+        if holding:
+            res = max(prices[i] + dp(i + 2, False), dp(i + 1, True))
+        else:
+            res = max(-prices[i] + dp(i + 1, 1), dp(i + 1, 0))
 
-    return dp(0, True)
+        return res
+
+    return dp(0, False)

dynamic_programming/test/test_max_profit_with_cooldown.py

@@ -1,14 +1,14 @@
 import unittest
 
-from dynamic_programming.BestTimeToBuyAndSellStockWithCooldown import max_profit
+from dynamic_programming.BestTimeToBuyAndSellStockWithCooldown import max_profit_top_down
 
 
 class MyTestCase(unittest.TestCase):
     def test_max_profit(self):
-        self.assertEqual(3, max_profit(self, [1, 2, 3, 0, 2]))
+        self.assertEqual(3, max_profit_top_down(self, [1, 2, 3, 0, 2]))
 
     def test_max_profit_1(self):
-        self.assertEqual(0, max_profit(self, [1]))
+        self.assertEqual(0, max_profit_top_down(self, [1]))
 
 
 if __name__ == '__main__':