Given an array A of 0s and 1s, we may change up to K values from 0 to 1.

Return the length of the longest (contiguous) subarray that contains only 1s.

Example 1:

Input: A = [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.

Naïve solution

The naïve solution is enumerate all the range pairs of array, check whether we have a consecutive sub-array to meet the required property.

The time complexity is: O(N^3).

class Solution {
public:
    int longestOnes(vector<int>& A, int K) {
        int ans = 0;
        for(int i=0; i<A.size(); i++) {
            for(int j=i; j<A.size(); j++) {
                int count = 0;
                for(int k=i; k<=j; k++) {
                    if(A[k] == 0) //change to 1
                        count++;
                    if(count <= K)
                        ans = max(ans, k-i+1);
                    else 
                        break;
                }
            }
        }
        return ans;
    }
};

Of course, this is a very slow solution, and we will get time-limit exceed.

Solution with approach of two-pointers

In the previous solution, there are duplicated computation. For example, [1~5] will contains the computation of [1~4], [1~3] etc.

An simpler solution is called two-pointers approach.

We define two variables calledl and r to store the left-index and right-index of a window. What we need to do is keep the required property and move l or r according to current status.

We can use a count to store the number of 0s in the window. When the number of 0s in the window is greater than K, we need to shrink the window. In every iteration, we compare the window size with the final result to determine whether to update the result.

Time complexity: O(n)
Space complexity: O(1)

Cpp

class Solution {
public:
    int longestOnes(vector<int>& A, int K) {
        int count = 0;
        int max_count = 0;
        int ans = 0;
        int l = 0;
        int r = 0;
        while(r < A.size()) { 
            int v = A[r];
            if(A[r] == 0) count++;
            while(count > K) {
                if(A[l] == 0)
                    count--;
                l++;
            }
            ans = max(ans, r - l + 1);
            r++;
        }
        return ans;
    }
};

Java

class Solution {

public int longestOnes(int[] A, int K) {
    int max = 0;
    for (int left = 0, right = 0; right < A.length; right++) {
        if (A[right] == 0) {
            if (K == 0) {
                while (A[left] == 1) {
                    left++;
                }
                left++;
            } else {
                K--;
            }
        }
        max = Math.max(right - left + 1, max);
    }
    return max;
}

};

Python

def longestOnes_concise(A, K):
    left,right,count = 0,0,0 
    for right in range(len(A)):
        if A[right] == 0:
            count += 1
        if count > K:
            if A[left] == 0:
                count -= 1
            left += 1
    return right-left+1

Preparing for an interview? Check out this!