是否有替代for循环的NumPy方案解决该数组匹配问题？

问题背景 #

我有三个长度相同的NumPy数组：

import numpy as np

weights = np.array([10, 14, 18, 22, 26, 30, 32, 34, 36, 38, 40])
resistances = np.array([15, 16.5, 18, 19.5, 21, 24, 27, 30, 33, 36, 39])
depths = np.array([0,1,2,3,4,5,6,7,8,9,10])

需求说明 #

遍历weights中的每个元素，在resistances中找到大于等于该元素的最近匹配项，通过其索引获取对应的depths值；若weights中的元素在resistances中无匹配项（即所有resistances元素都小于该值），则返回depths的最后一个值，最终将结果存入列表。

原for循环实现 #

目前已实现的for循环版本如下：

SWP = []
for w in weights:
    if len(depths[w<=resistances]) == 0:
        swp=depths[-1]
    else:
        swp = np.min(depths[w<=resistances])
    SWP.append(swp)
SWP

无循环的NumPy优化方案 #

可以利用np.searchsorted实现向量化操作，彻底避免循环，提升运行效率：

import numpy as np

weights = np.array([10, 14, 18, 22, 26, 30, 32, 34, 36, 38, 40])
resistances = np.array([15, 16.5, 18, 19.5, 21, 24, 27, 30, 33, 36, 39])
depths = np.array([0,1,2,3,4,5,6,7,8,9,10])

# 找到每个weight对应的第一个>=它的resistances元素索引
indices = np.searchsorted(resistances, weights, side='left')
# 处理超出resistances范围的情况，将索引设为最后一个元素的位置
indices[indices >= len(resistances)] = len(resistances) - 1
# 根据索引提取对应depth值并转为列表
SWP = depths[indices].tolist()
print(SWP)

方案说明 #

• np.searchsorted(resistances, weights, side='left')：因为resistances是升序排列的，这个函数可以直接为每个weights元素找到第一个大于等于它的元素索引，完美匹配需求里的「最近匹配项」逻辑。
• 当weights元素大于resistances最大值时，searchsorted会返回等于数组长度的索引，这时候我们将其替换为数组最后一个元素的索引，以此获取depths[-1]。
• 整个过程是纯向量化操作，没有循环逻辑，数组规模越大，相比原for循环版本的效率提升越明显。

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