如何用Numpy完全向量化实现支持间隔容忍的二进制序列连续1区间提取

优化后的完全向量化Numpy实现 #

原有实现中的for循环可以通过Numpy的repeat广播操作完全替换，无需显式迭代，长序列场景下性能提升明显，逻辑也更简洁：

import numpy as np

def get_segments(segment_of_sums, mingap=1):
    nonempty_padded = np.concatenate(([0], segment_of_sums, [0]))
    b = nonempty_padded > 0
    edges = b[:-1] ^ b[1:]
    indices = np.argwhere(edges)[:, 0]
    
    if mingap > 1 and len(indices) > 2:
        gap_lengths = indices[2::2] - indices[1::2][:-1]
        gap_keeps = gap_lengths >= mingap
        index_keeps = np.zeros_like(indices, dtype=np.bool_)
        index_keeps[[0, -1]] = True
        # 向量化替换原有for循环：每个保留间隔对应前后两个边界都保留
        index_keeps[1:-1] = gap_keeps.repeat(2)
        indices = indices[index_keeps]
    
    return indices.reshape((-1, 2))

# 测试代码
if __name__ == "__main__":
    for L in [
        [0,1,0, 1,1,1, 0,0,1, 1,1,0],
        [0,1,0, 1,1,1, 0,0,1, 1,1,1],
        [1,1,0, 1,1,1, 0,0,1, 1,1,1]
    ]:
        A = np.array(L, dtype=np.bool_)
        print(get_segments(A, mingap=2))

优化逻辑说明 #

原有for循环的作用是为每个符合保留条件的间隔，标记对应的前后两个边界索引为保留状态。我们可以利用np.repeat将每个gap_keeps的布尔值直接复制两份，刚好对应间隔前后两个边界的标记位置，直接批量赋值即可完成原来for循环的所有操作，全程无Python级别的循环，完全运行在Numpy的向量化算子上，性能随序列长度增长近似线性。

测试输出和原有实现完全一致：

[[ 1  6]
 [ 8 11]]
[[ 1  6]
 [ 8 12]]
[[ 0  6]
 [ 8 12]]

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