查找数组 B 中与数组 A 中最匹配的元素的索引

Question

我有两个数组A和B。现在让它们都是一维的。
对于 中的每个元素，我需要与 中的元素最匹配的元素A的索引。BA

我可以使用列表表达式来解决这个问题

import numpy as np

A = np.array([ 1, 3, 1, 5 ])
B = np.array([ 1.1, 2.1, 3.1, 4.1, 5.1, 6.1 ])

indices = np.array([ np.argmin(np.abs(B-a)) for a in A ])

print(indices)    # prints [0 2 0 4]
print(B[indices]) # prints [1.1 3.1 1.1 5.1]

但这种方法对于巨大的数组来说确实很慢。
我想知道是否有更快的方法利用优化的 numpy 函数。

Answer 1

您可以计算 B 和重塑 A 之间的绝对差，并使用argminon axis=1：

np.argmin(np.abs(B-A[:,None]), axis=1)

输出：array([0, 2, 0, 4])