基于已排序numpy数组对另一数组元素向下取最近值赋值的高效方法

最优实现方案 #

你可以直接使用numpy内置的np.searchsorted函数实现，该函数基于二分查找实现，完全符合numpy向量化编程规范，时间复杂度远低于你当前的循环实现，尤其适合数据量较大的场景。

实现代码 #

import numpy as np

arr1 = np.array([0, 3, 7, 9, 13])
arr2 = np.array([1, 2, 7, 8, 12, 14])

# 计算插入位置后减1，得到arr1中小于等于对应元素的最大索引
idx = np.searchsorted(arr1, arr2, side='right') - 1
arr3 = arr1[idx]

print(arr3)

输出结果 #

array([ 0,  0,  7,  7,  9, 13])

原理解释 #

1. np.searchsorted默认接收已排序的一维数组作为第一个参数，返回第二个参数中每个元素插入到第一个数组后仍然保持排序的位置索引
2. 指定side='right'时，返回的是相同元素的右侧插入位置，减1后正好对应「小于等于目标值的最大元素」的索引
3. 如果你的场景中存在arr2元素小于arr1最小值的情况，计算得到的索引会是-1，你可以根据需求额外增加边界处理逻辑，比如将这些位置赋值为默认值。

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