求O(n)时间复杂度的0-1数组特定子数组计数算法

首先通过数学推导简化原问题的核心条件：

定义变量 #

• 设数组总长度为 n，数组中1的总个数为 total_1，则0的总个数 total_0 = n - total_1。
• 对于任意子数组 [L, R]（索引从L到R，闭区间），记其长度为 len = R - L + 1，子数组内1的个数为 cnt1，子数组内0的个数为 len - cnt1。

条件转化 #

原要求是：cnt1 > 子数组外0的数量
子数组外0的数量 = 总0数 - 子数组内0数 = total_0 - (len - cnt1)

代入原条件并化简：

cnt1 > total_0 - (len - cnt1)
cnt1 > total_0 - len + cnt1
0 > total_0 - len
len > total_0

结论：原条件等价于「子数组的长度大于数组中0的总个数」

计算满足条件的子数组个数 #

总子数组个数为 n*(n+1)/2，减去长度≤total_0的子数组个数即可得到答案：

1. 遍历数组计算 total_0（O(n)时间）
2. 计算长度≤total_0的子数组个数：
   • 如果 n ≤ total_0：所有子数组都不满足条件，答案为0
   • 否则：长度为1的子数组有n个，长度为2的有n-1个，…，长度为total_0的有n - total_0 + 1个，总和为 total_0*(2*n - total_0 + 1)/2
3. 最终答案 = 总子数组数 - 长度≤total_0的子数组数

结合你的示例验证 #

你的输入数组：[1, 0, 0, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1]

• n=13，total_0=6（数出数组中0的个数为6）
• 总子数组数：13*14/2=91
• 长度≤6的子数组数：6*(2*13 -6 +1)/2 = 63
• 答案：91-63=28

代码实现（O(n)时间） #

def count_valid_subarrays(arr):
    n = len(arr)
    total_0 = arr.count(0)
    total_subarrays = n * (n + 1) // 2
    if n <= total_0:
        return 0
    k = total_0
    valid_less = k * (2 * n - k + 1) // 2
    return total_subarrays - valid_less

# 测试示例
arr = [1, 0, 0, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1]
print(count_valid_subarrays(arr))  # 输出28

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