如何查找Numpy数组中True值对应连续索引的起止位置

你可以用NumPy的差分运算快速完成需求，下面分两种场景给出实现代码： #

场景1：从已提取的True索引数组计算 #

原理 #

连续的索引相邻差值为1，只要找到差值大于1的位置，就能定位连续区间的分割点，再分别提取起止即可。

代码示例 #

import numpy as np

# 你的True索引数组
idx = np.array([20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36,
       37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 60, 61, 62, 63,
       64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80,
       81, 82, 83, 84, 85, 86, 87, 88, 89, 90])

# 计算相邻索引的差值
diff_arr = np.diff(idx)
# 定位连续区间的分割点
split_pos = np.where(diff_arr > 1)[0]

# 提取起始和结束索引
start_indices = np.hstack([[idx[0]], idx[split_pos + 1]])
end_indices = np.hstack([idx[split_pos], [idx[-1]]])

# 输出结果
print("start_indices =", start_indices)
print("end_indices =", end_indices)

运行结果 #

start_indices = [20 60]
end_indices = [49 90]

场景2：直接从原始布尔数组计算 #

如果还没提取True索引，可以直接从布尔数组计算，步骤更简洁：

代码示例 #

import numpy as np

# 你的原始布尔数组，替换成你实际的mask即可
mask = np.array([False]*20 + [True]*30 + [False]*10 + [True]*31 + [False]*119)

# 给布尔数组前后各补一个False，处理边界的连续区间
padded_mask = np.pad(mask, (1, 1), constant_values=False)
# 计算差分，从False变True差值为1，对应区间起始；True变False差值为-1，对应区间结束的下一位
mask_diff = np.diff(padded_mask)

start_indices = np.where(mask_diff == 1)[0]
end_indices = np.where(mask_diff == -1)[0] - 1

# 输出结果
print("start_indices =", start_indices)
print("end_indices =", end_indices)

运行结果 #

start_indices = [20 60]
end_indices = [49 90]

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