如何高效获取Numpy布尔数组中首个和末尾True值的索引？

问题描述 #

我有一个大型Numpy布尔数组，其中True代表另一预处理信号的过零条件。我只需要数组中首个和末尾True值的索引，不想使用np.where或np.argwhere（避免不必要地遍历整个数组），也不愿将数组转换为list并使用index()方法。目前自己实现的方案如下：

# Alias for compactness
enum = enumerate

# Example array, desired output = (3, -5)
a = np.array([False, False, False, True, ..., True, False, False, False, False])

out = next(idx for idx in zip((i for i, x in enum(a) if x), (-i-1 for i, x in enum(a[::-1]) if x)))

print(out)
# (3, -5)

但觉得这个方案略显繁琐，想知道Numpy是否有原生语法能实现这类功能？

解决方案 #

1. 简化惰性生成器写法 #

你的惰性求值思路本身是高效的（找到目标后立即停止遍历），可以简化代码提升可读性：

import numpy as np

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

# 获取第一个True的正索引
first_idx = next(i for i, val in enumerate(a) if val)
# 从后往前找第一个True，转换为原数组的负索引
last_neg_idx = next(-(i + 1) for i, val in enumerate(reversed(a)) if val)

out = (first_idx, last_neg_idx)
print(out)  # 输出 (3, -5)

2. Numpy原生简洁写法（需遍历整个数组） #

如果不在意遍历整个数组（Numpy向量化操作在大型数组上仍有不错性能），可以用argmax快速实现：

import numpy as np

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

if a.any():
    first_idx = a.argmax()  # 返回数组中第一个True的索引
    # 反转数组后找到第一个True，再转换为原数组的负索引
    last_neg_idx = - (a[::-1].argmax() + 1)
    out = (first_idx, last_neg_idx)
else:
    # 处理数组中无True的边界情况
    out = (None, None)

print(out)  # 输出 (3, -5)

3. 关于Numpy原生短路查找 #

目前Numpy没有原生的短路式查找方法（即找到第一个匹配项就停止遍历），因为其核心设计是向量化批量处理，通常会遍历整个数组。因此，你的惰性生成器思路在大型数组场景下是更高效的选择，简化后的写法已经足够简洁清晰。

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