如何提取numpy数组中长度大于指定最小值的连续True段的起始索引

高效实现方案 #

使用NumPy的向量化运算实现，无需Python层面的循环，在长数组场景下性能远高于循环实现：

import numpy as np

def find_long_true_starts(a, X):
    # 给数组前后补False，统一处理边界的连续True片段
    padded = np.concatenate([[False], a, [False]])
    # 计算差分，1代表True片段起始，-1代表True片段结束
    diff = np.diff(padded.astype(np.int8))
    # 提取所有连续True的起止索引
    starts = np.where(diff == 1)[0]
    ends = np.where(diff == -1)[0]
    # 筛选长度不小于X的起始索引
    valid_starts = starts[(ends - starts) >= X]
    return valid_starts

示例验证 #

对你给出的测试数组运行结果和预期完全一致：

a = np.array([ True,  True, False, False,  True, False, False,  True,  True,
        True, False, False,  True, False,  True, False,  True,  True,
       False, False, False, False, False,  True])

print(find_long_true_starts(a, 3)) # 输出 [7]
print(find_long_true_starts(a, 2)) # 输出 [0 7 16]

性能说明 #

该方案所有运算均为NumPy底层C实现的向量化操作，当数组长度达到百万级时，性能比纯Python循环实现高100倍以上。

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