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

solairerove · Nov 22, 2023 · bba5a1a · bba5a1a
1 parent 5c852df
commit bba5a1a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/dynamic_programming/BestTimeToBuyAndSellStockWithCooldown.py b/dynamic_programming/BestTimeToBuyAndSellStockWithCooldown.py
@@ -2,6 +2,19 @@
 from typing import List
 
 
+# O(n) time || O(1) space
+def max_profit(self, prices: List[int]) -> int:
+    buy, sell, cooldown = -prices[0], 0, 0
+    for price in prices[1:]:
+        _buy = max(buy, cooldown - price)
+        _sell = buy + price
+        _cooldown = max(cooldown, sell)
+
+        buy, sell, cooldown = _buy, _sell, _cooldown
+
+    return max(sell, cooldown)
+
+
 # O(n) time || O(n) space
 def max_profit_top_down(self, prices: List[int]) -> int:
     @lru_cache(None)
@@ -12,7 +25,7 @@ def dp(i, holding):
         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))
+            res = max(-prices[i] + dp(i + 1, True), dp(i + 1, False))
 
         return res
 

diff --git a/dynamic_programming/test/test_max_profit_with_cooldown.py b/dynamic_programming/test/test_max_profit_with_cooldown.py
@@ -1,13 +1,15 @@
 import unittest
 
-from dynamic_programming.BestTimeToBuyAndSellStockWithCooldown import max_profit_top_down
+from dynamic_programming.BestTimeToBuyAndSellStockWithCooldown import max_profit_top_down, max_profit
 
 
 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]))