如何简洁获取numpy一维数组首个连续True序列的首尾索引

问题描述 #

我有一个numpy一维布尔数组，示例如下：

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

需要获取两个索引：

• 数组中第一个True值的索引（此示例为1）
• 首个连续True序列的最后一个索引（此示例为3）

假设数组中至少存在一个True值，以下是更多测试示例及预期输出：

[False, True, True] # 预期输出：1, 2
[True, True, True, False, True, True, False] # 预期输出：0, 2
[False, True, True, True] # 预期输出：1, 3
[False, True, False, True] # 预期输出：1, 1

简洁实现方案 #

利用numpy内置的高效函数可以快速实现需求，这里提供两种简洁的写法：

方法一：基于np.argmax的直接定位

import numpy as np

def get_first_true_range(arr):
    # 找到第一个True的索引（True对应数值1，argmax返回第一个最大值的位置）
    start = np.argmax(arr)
    # 从start位置开始，找第一个False的索引
    first_false_after_start = start + np.argmax(arr[start:] == False)
    # 如果从start开始全是True，修正为数组最后一个索引
    if arr[start:].all():
        end = len(arr) - 1
    else:
        end = first_false_after_start - 1
    return start, end

方法二：基于np.diff的边界检测

通过检测布尔值变化的边界来定位连续序列的终点：

import numpy as np

def get_first_true_range(arr):
    start = np.argmax(arr)
    # 将布尔数组转为int后计算差分，-1表示从True变为False
    boundaries = np.diff(arr.astype(int))
    # 查找start之后第一个边界为-1的位置
    end_indices = np.where(boundaries[start:] == -1)[0]
    if len(end_indices) == 0:
        end = len(arr) - 1
    else:
        end = start + end_indices[0]
    return start, end

测试验证 #

将上述函数应用到示例数组中，均可得到预期结果：

# 测试示例1
a1 = np.array([False, True, True, True, False, True, True, False])
print(get_first_true_range(a1))  # 输出 (1, 3)

# 测试示例2
a2 = np.array([False, True, True])
print(get_first_true_range(a2))  # 输出 (1, 2)

# 测试示例3
a3 = np.array([True, True, True, False, True, True, False])
print(get_first_true_range(a3))  # 输出 (0, 2)

# 测试示例4
a4 = np.array([False, True, True, True])
print(get_first_true_range(a4))  # 输出 (1, 3)

# 测试示例5
a5 = np.array([False, True, False, True])
print(get_first_true_range(a5))  # 输出 (1, 1)

内容的提问来源于stack exchange，提问作者Salvatore Daniele Bianco