二进制数组含k个1的最短子数组O(n)算法实现及代码错误排查

原代码存在的问题 #

• 拼写错误：代码中更新结果时错把result写成了resut，导致结果永远无法正确更新
• 边界判断错误：数组索引最大值为len(arr)-1，原代码中ep>n的判断会触发数组越界异常
• 滑动窗口逻辑错误：没有维护窗口内1的计数，每次移动指针的逻辑无法覆盖所有可能的有效窗口，会漏掉更短的子数组情况
• 异常场景未处理：如果数组中1的总数量小于k，原代码会直接返回错误的初始值(0,0)

正确算法思路（时间复杂度O(n)） #

采用双指针滑动窗口方案：

1. 初始化左右指针left=0，窗口内1的计数count_one=0，最短子数组的起止位置初始值设为无效值
2. 右指针遍历整个数组，遇到1则count_one +=1
3. 当count_one ==k时，尝试向右移动左指针缩小窗口，每次移动前计算当前窗口长度，若小于已记录的最短长度则更新结果，直到窗口内1的数量小于k为止
4. 遍历结束后返回最短子数组的起止位置即可

修复后可运行代码 #

def subarray(arr, k):
    n = len(arr)
    left = 0
    count_one = 0
    min_len = float('inf')
    result = (-1, -1)  # 初始无效值，表示未找到符合要求的子数组
    
    for right in range(n):
        if arr[right] == 1:
            count_one += 1
        
        # 当窗口内1的数量达到k时，尝试缩小左边界找更短的子数组
        while count_one == k:
            current_len = right - left
            # 更新最短子数组记录
            if current_len < min_len:
                min_len = current_len
                result = (left, right)
            # 移动左指针
            if arr[left] == 1:
                count_one -= 1
            left += 1
    
    return result if result[0] != -1 else None

# 测试用例
A=[0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1,
1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 0, 0, 1, 1, 0]
k = 7
print(subarray(A,k))

测试结果 #

运行上述代码，输入示例的输出为(9, 16)，对应子数组长度为7，是包含7个1的最短子数组。

内容的提问来源于stack exchange，提问作者Sahand Shabanlouei