[LeetCode] 1004. Max Consecutive Ones III

1. Question

Given a binary array nums and an integer k, return the maximum number of consecutive 1's in the array if you can flip at most k 0's.

Example 1:

Input: nums = [1,1,1,0,0,0,1,1,1,1,0], k = 2
Output: 6
Explanation: [1,1,1,0,0,1,1,1,1,1,1]
Bolded numbers were flipped from 0 to 1. The longest subarray is underlined.

Example 2:

Input: nums = [0,0,1,1,0,0,1,1,1,0,1,1,0,0,0,1,1,1,1], k = 3
Output: 10
Explanation: [0,0,1,1,1,1,1,1,1,1,1,1,0,0,0,1,1,1,1]
Bolded numbers were flipped from 0 to 1. The longest subarray is underlined.

2. Thoughts

3. Tips learned

4. My solution

class Solution:
    def longestOnes(self, nums: List[int], k: int) -> int:
        final_max = current_max = index = taken = zero_index = 0
        zeros = []
        while index < len(nums):
            # if 1
            if nums[index] == 1:
                current_max += 1
            # if 0
            else:
                if taken == k:
                    zeros.append(index)
                    final_max = max(final_max, current_max)
                    current_max = index - (zeros[zero_index] + 1) + 1
                    zero_index += 1
                else:
                    while index < len(nums) and nums[index] == 0 and taken < k:
                        zeros.append(index)
                        current_max += 1
                        index += 1
                        taken += 1
                    index -= 1
            index += 1

        return max(final_max, current_max)

55/55 cases passed (492 ms)
Your runtime beats 77.71 % of python3 submissions
Your memory usage beats 6.4 % of python3 submissions (17.5 MB)

5. AI Solution and Improvements

Your algorithm is correct but could be simplified.

The idea is to maintain two pointers to mark the boundaries of your current subarray, expanding the right end until you've added k zeroes. If you get a zero and you've already added k zeroes, you shrink the subarray from the left end until a zero is removed.

Here is the refactored code:

class Solution:
    def longestOnes(self, nums: List[int], k: int) -> int:
        max_length = 0
        left = 0
        zeros = 0  # Count the number of zeros in the current window
        
        for right in range(len(nums)):
            if nums[right] == 0:
                zeros += 1
            
            # Shrink window from the left if we have more than k zeros
            while zeros > k:
                if nums[left] == 0:
                    zeros -= 1
                left += 1
            
            max_length = max(max_length, right - left + 1)
        
        return max_length

55/55 cases passed (492 ms)
Your runtime beats 77.71 % of python3 submissions
Your memory usage beats 6.4 % of python3 submissions (17.5 MB)

This refactored code should be easier to understand and debug, and it has the same time complexity as your original solution.