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

【问题标题】：Find indices for elements in array B best matching those in array A查找数组 B 中与数组 A 中的元素最匹配的元素的索引
【发布时间】：2021-11-03 14:09:15
【问题描述】：

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

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

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 函数。

【问题讨论】：

标签： python arrays performance numpy indexing

【解决方案1】：

您可以计算 B 和重构 A 之间的绝对差，并在 axis=1 上使用 argmin：

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

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

【讨论】：

啊，广播救援，太棒了！谢谢！
@Bastian 请注意，中间数组的大小将是两个输入的乘积，这可能会占用大量内存，具体取决于输入和资源；）

【解决方案2】：

广播可能会咬你一口（tmp数组的创建也会及时包含在内），下面的方法不使用大量的tmp内存，所以内存效率很高。 Here 是由于内存使用过多导致广播速度变慢时的参考

保留此内容以供参考。除此之外，您可以在 cython numpy 中编写自定义函数。与 numba 相比，Cython 使用不同的优化。所以需要试验哪一种优化更好。但是对于 numba，您可以留在 python 中并编写类似 c 的代码

import numpy as np
import numba as nb

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

# Convert to fast optimized machine code
@nb.njit(
    # Signature of Input
    (nb.float64[:], nb.float64[:]),
    # Optional
    parallel=True
)
def less_mem_ver(A, B):

    arg_mins = np.empty(A.shape, dtype=np.int64)

    # nb.prange is for parallel=True
    # what can be parallelized
    # Usage of for loop because, to prevent creation of tmp arrays due to broadcasting
    # It takes time to allocate tmp array
    # No loss in writing for loop as numba will vectorize this just like numpy
    for i in nb.prange(A.shape[0]):
        min_num = 1e+307
        min_index = -1
        for j in range(B.shape[0]):
            t = np.abs(A[i] - B[j])
            if t < min_num:
                min_index = j
                min_num = t
        arg_mins[i] = min_index
    return arg_mins
less_mem_ver(A, B)

【讨论】：

【解决方案3】：

除了已经使用广播给出的答案外，还有一个是内部广播。

我想将我的其他答案分开，因为它使用 numpy 以外的其他答案

np.argmin(np.abs(np.subtract.outer(A, B)), axis=1)

您可以在许多 ufunc 上调用 outer。

Reference

【讨论】：