布尔序列中True到下一个True的False计数（向量化实现）

针对布尔序列的处理需求，这里提供一个完全基于NumPy的向量化实现，无需显式迭代，效率远高于反向循环：

import numpy as np

def compute_true_interval(arr):
    # 转换为NumPy布尔数组
    bool_arr = np.asarray(arr, dtype=bool)
    # 获取所有True值的索引位置
    true_positions = np.where(bool_arr)[0]
    result = np.full(len(bool_arr), np.nan)
    
    if not len(true_positions):
        return result
    
    # 计算相邻True之间的False元素数量：后一个True索引 - 当前True索引 - 1
    intervals = np.diff(true_positions) - 1
    # 最后一个True后续无True，填充np.inf
    intervals = np.append(intervals, np.inf)
    
    # 将计算结果映射回原数组的True位置
    result[true_positions] = intervals
    
    return result

测试示例 #

input_seq = [False, False, True, True, False, False, True, False, True, False , False, False]
output = compute_true_interval(input_seq)
print(output)

输出结果：

array([nan, nan,  0.,  2., nan, nan,  1., nan, inf, nan, nan, nan])

实现思路解析 #

• 向量化提取True索引：用np.where一次性获取所有True值的位置，避免逐个遍历判断。
• 批量计算间隔：np.diff直接计算相邻True索引的差值，减1后得到两个True之间的False元素个数，这一步是完全的向量化运算，底层由C优化实现。
• 批量赋值：通过索引直接将计算结果填充到结果数组的对应位置，替代循环赋值，大幅提升处理大数组的效率。

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