避免 NumPy 1D 最近邻中的 for 循环

Question

我有以下代码，其中我得到一维中的 N 个最近邻居：

import numpy as np

def find_nnearest(arr, val, N):
    idxs = []
    for v in val:
        idx = np.abs(arr - v).argsort()[:N] 
        idxs.append(idx)
    return np.array(idxs)

A = np.arange(10, 20)
test = find_nnearest(A, A, 3)
print(test)

明显使用了 for 循环来获取

idx

。有没有一种numpythonic方法来避免这个for循环（但返回相同的数组）？

Answer 1

要回答你的问题，是的，你可以通过避免显式的 for 循环来使代码更加“numpythonic”。您可以使用广播在单个操作中计算整个数组 arr 和数组 val 之间的绝对差异。这是代码的更新版本：

import numpy as np

def find_nnearest(arr, val, N):
    idxs = np.abs(arr[:, None] - val).argsort(axis=0)[:N]
    return idxs

A = np.arange(10, 20)
test = find_nnearest(A, A, 3)
print(test)

在此版本中，arr[:, None] 为 arr 引入了一个新轴，使其成为一个二维数组，其中每一行对应 arr 中的一个元素。然后，广播负责在一个步骤中计算 arr 和 val 之间所有元素组合的绝对差。其余代码保持类似，但现在使用 axis=0 完成索引，指定应沿新轴进行排序。