Updated 1004

Author: trung-hn

src/1004.max-consecutive-ones-iii.py

@@ -6,46 +6,39 @@
 
 # @lc code=start
 # TAGS: Sliding Window, Two Pointers
+from typing import List
+
+
 class Solution:
-    # 6465 ms, 5.02%. Time: O(N^2). Space: O(N)
-    def longestOnes(self, A: List[int], K: int) -> int:
-        
-        def sum_from_i(i, K):
-            k = K
-            rv = 0
-            for val, length in group[i:]:
-                if val == 1:
-                    rv += length
-                elif val == 0 and length <= k:
-                    k -= length
-                    rv += length
-                elif val == 0 and length > k:
-                    rv += k
-                    k = 0
-                    break
-            return rv, k
-        
-        group = [(i, len(list(g))) for i, g in itertools.groupby(A)]
-        
-        ans = 0
-        for i in range(len(group)):
-            rv, k = sum_from_i(i, K)
-            if i > 0 and group[i - 1][1] > k:
-                rv += k
-            ans = max(ans, rv)
-            
+    # 652 ms, 41.04%
+    # Easy to understand solition
+    def longestOnes(self, A: List[int], k: int) -> int:
+        ptr = ans = 0
+        for i, n in enumerate(A):
+            if not n:
+                while not k:
+                    k += A[ptr] == 0
+                    ptr += 1
+                k -= 1
+            ans = max(ans, i + 1 - ptr)
         return ans
-    # 552 ms, 94.19%. Time: O(N). Space: O(1). Sliding Window
+
+    # 564 ms, 86.48%
+    # Over simplified of the version above. Very clever but hard to understand
+    # We never decrease the size of the window
+    # K is used as the offset.
+    # if K is positive, we can always increase the window size
+    # if K is negative, we increase K (if applicable) and this it as offset at the end
+    # at some point, K might turn back to positive again, at which time the offset reduce to 0
     def longestOnes(self, A: List[int], K: int) -> int:
         left = 0
         for right in range(len(A)):
-            K -= 1 - A[right]
+            K -= A[right] == 0
+
             if K < 0:
-                K += 1 - A[left]
+                K += A[left] == 0
                 left += 1
-                
-        # we can return this because we never decrease the size, only increase of not change.
+
         return right - left + 1
 
 # @lc code=end
-