如何在三维NumPy数组中高效应用searchsorted匹配对应值？

给定三维NumPy数组a（形状(1000,2000,200)）和二维数组b（形状(1000,2000)），两者均无NaN值：

import numpy as np
a = np.random.randint(0,50,size=(1000,2000,200))
b = np.random.randint(0,50,size=(1000,2000,))

需求是找到每个a[i,j]子数组中最接近b[i,j]值的元素索引。

最初尝试用np.searchsorted：

indices = np.searchsorted(a, b)

但该函数仅支持一维数组，运行报错ValueError: object too deep for desired array。

已尝试一种基于差分定位的方法：

b_3d = np.repeat(b[:, :, np.newaxis], a.shape[2], axis=2)
diff = np.abs(a - b_3d)
indices = np.argmin(diff, axis=2)

现寻求更高效的向量化方法替代非向量化实现。

高效向量化优化方案 #

你当前的差分方法逻辑可行，但可以优化内存占用和计算效率——不需要用np.repeat扩展b的维度，利用NumPy的广播机制就能直接计算差值：

# 利用广播直接计算每个a[i,j,k]与b[i,j]的绝对差
diff = np.abs(a - b[..., np.newaxis])
# 沿第三个维度找最小差的索引
indices = np.argmin(diff, axis=2)

优化点说明 #

• 广播机制避免了显式复制b的数据，b[..., np.newaxis]只是给b增加一个长度为1的第三维度，内存占用远低于np.repeat生成的完整三维数组
• 计算逻辑和原方法一致，但内存效率提升显著，尤其当a.shape[2]（此例为200）较大时，节省的内存非常可观

进一步优化（若a[i,j]为有序数组） #

如果每个a[i,j]子数组是升序排列的，可结合np.searchsorted实现更快的二分查找：

# 为每个a[i,j]子数组执行searchsorted，找到b[i,j]的插入位置
pos = np.apply_along_axis(np.searchsorted, 2, a, b[..., np.newaxis])
# 处理边界：插入位置为0时选索引0；等于子数组长度时选最后一个索引
pos = np.clip(pos, 1, a.shape[2]-1)
# 比较插入位置前后元素与b[i,j]的差值
left_diff = np.abs(a[np.arange(a.shape[0])[:, None], np.arange(a.shape[1]), pos-1] - b)
right_diff = np.abs(a[np.arange(a.shape[0])[:, None], np.arange(a.shape[1]), pos] - b)
# 选择差值更小的索引（相等时选靠前的）
indices = np.where(left_diff <= right_diff, pos-1, pos)

这种方法时间复杂度更低（二分查找 vs 遍历所有元素），但前提是每个a[i,j]必须有序。

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