【leetcode】1004. Max Consecutive Ones III

题目如下：

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.

Example 2:

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

 

Note:

1. 1 <= A.length <= 20000
2. 0 <= K <= A.length
3. A[i] is 0 or 1 

解题思路：如果我们把所有0所在位置对应的下标存入一个数组inx_0中，例如Example 2的inx_0 = [0,1,4,5,9,12,13,14]，因为最多把K个0变成1，要构造出最长的全为1的连续子数组，那么很显然所有进行反转操作的0所对应下标在inx_0中必须是连续的。接下来开始遍历数组A，在K大于0的条件下，遇到1则count_1加1，遇到0的话则K减1并且count_1加1（表示这个0反转成1）；如果在K=0的情况下遇到0，那么需要去掉第一个0变成的1和这个0之前连续的1的数量，即count减1（减去第一个0变成1的计数），再减去第一个0前面连续的1的数量i，记为count_1减i。记录count_1出现的最大值，直到数组A遍历完成为止。

代码如下：

class Solution(object):
    def longestOnes(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: int
        """
        count_1 = 0
        zero_inx = []
        res = 0
        zero_before = -1
        for i in range(len(A)):
            if A[i] == 1:
                count_1 += 1
            else:
                zero_inx.append(i)
                if K > 0:
                    K -= 1
                    count_1 += 1
                elif K == 0:
                    res = max(res,count_1)
                    first_0 = zero_inx.pop(0)
                    count_1 -= (first_0 - zero_before - 1)
                    zero_before = first_0
        res = max(res, count_1)
        return res