使用欧几里德距离在numpy数组列表中查找numpy数组的最近邻居

Question

我有一个n维向量，我想在使用欧几里德距离的n维向量列表中找到它的k个最近邻居。

我编写了以下代码（k = 10），但是运行得太慢了，我想知道是否有更优化的解决方案。

def nearest_neighbors(value, array, nbr_neighbors=1):
    return np.argsort(np.array([np.linalg.norm(value-x) for x in array]))[:nbr_neighbors]

Answer 1

使用scipy＆＃39; kd-tree。

一个小例子是available here。

许多人似乎抱怨性能并建议sklearn's implementation（链接sklearn.neighbors，内部使用此data-structure）！

Answer 2

正如sascha所说，我最终使用了scipy库（但NearestNeighbors方法），将计算时间从50小时缩短到36分钟。这是一种我不应该尝试重新实现自己的计算，因为专用库对此更加优化。

NearestNeighbors方法还允许您传入值列表并返回每个值的k个最近邻居。

最终代码是：

def nearest_neighbors(values, all_values, nbr_neighbors=10):
    nn = NearestNeighbors(nbr_neighbors, metric='cosine', algorithm='brute').fit(all_values)
    dists, idxs = nn.kneighbors(values)