Python中无需循环实现输入数组元素匹配上下限参考数组的索引定位

要在不使用循环的前提下完成这个区间匹配任务，我们可以利用NumPy的矢量化操作，这里提供两种高效的实现方式：

方法一：利用np.searchsorted（推荐，高效且简洁） #

如果你的lower数组是严格递增的，并且upper[i]恰好等于lower[i+1]（即区间连续无重叠），那么np.searchsorted是最优选择：

import numpy as np

a = np.array([120,350,410,354,247])
lower = np.array([100,150,200,250,300,350,400,450])
upper = np.array([150,200,250,300,350,400,450,500])

# 计算每个元素对应的区间索引
output = np.searchsorted(lower, a, side='right') - 1
print(output)  # 输出：[0 5 6 5 2]

原理说明： #

• np.searchsorted(lower, a, side='right')会返回每个a元素在lower数组中右侧插入的位置——也就是第一个大于该a元素的lower值的索引。
• 减去1之后，就得到了满足lower[i] <= a < upper[i]的区间索引：
  • 比如a=350，lower中第一个大于350的是400（索引6），减1后得到5，正好对应lower[5]<=350<upper[5]的区间。
  • 这种方法完全矢量化，没有循环，处理大规模数组时性能极佳。

方法二：利用广播（适用于非连续区间场景） #

如果你的区间不是连续的，或者lower数组无序，那么可以用广播生成布尔掩码来匹配区间：

import numpy as np

a = np.array([120,350,410,354,247])
lower = np.array([100,150,200,250,300,350,400,450])
upper = np.array([150,200,250,300,350,400,450,500])

# 生成二维布尔掩码，mask[i,j]表示a[i]是否在第j个区间内
mask = (lower <= a[:, np.newaxis]) & (a[:, np.newaxis] < upper)
# 找到每行第一个为True的索引（即匹配的区间索引）
output = np.argmax(mask, axis=1)
print(output)  # 输出：[0 5 6 5 2]

原理说明： #

• a[:, np.newaxis]将一维数组a转换为列向量，和lower/upper广播成二维数组，实现每个a元素与所有区间的逐一比较。
• mask是一个布尔数组，每行对应一个a元素，每列对应一个区间，值为True表示该元素属于对应区间。
• np.argmax(mask, axis=1)返回每行中第一个True的索引，也就是该元素匹配的区间索引（如果有多个区间匹配，会返回第一个匹配的）。

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