如何提取NumPy数组中指定范围内每组值的首个元素

需求说明 #

给定NumPy数组和一个指定范围值，返回新数组，其中包含每个元素间差值不超过指定范围、且组内至少有2个元素的值组的首个元素。

示例 #

• 指定范围：specified_range = 10
• 输入数组：in_array = np.array([3, 14, 25, 37, 38, 50, 61, 66, 70])
• 符合条件的组：[37, 38]、[61, 66, 70]
• 期望结果：array([37, 61])

解决方案代码 #

import numpy as np

specified_range = 10
in_array = np.array([3, 14, 25, 37, 38, 50, 61, 66, 70])

# 计算相邻元素的差值，生成符合范围的掩码
diff = np.diff(in_array)
mask = diff <= specified_range

# 定位连续符合条件的块的边界
changes = np.diff(np.concatenate([[False], mask, [False]]))
block_starts = np.where(changes == 1)[0]
block_ends = np.where(changes == -1)[0]

# 筛选出组内元素≥2的块，提取对应起始元素
valid_blocks = (block_ends - block_starts) >= 1
result = in_array[block_starts[valid_blocks]]

print(result)  # 输出：[37 61]

步骤解释 #

1. 生成差值掩码：计算数组相邻元素的差值，用布尔值标记差值在指定范围内的位置。
2. 定位连续块边界：通过扩展掩码并计算变化点，找到所有连续符合条件的块的起始和结束索引。
3. 筛选有效块并提取结果：只保留对应原数组元素数≥2的块，提取这些块在原数组中的首个元素。

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