在x和y坐标的numpy数组中查找最近点的索引

Question

我有两个2d numpy数组：x_array包含x方向上的位置信息，y_array包含y方向上的位置。

然后，我有一长串x，y点。

对于列表中的每个点，我需要找到最接近该点的位置（在数组中指定）的数组索引。

我已经根据这个问题天真的产生了一些有效的代码：
在numpy数组中找到最接近的值

即

import time
import numpy

def find_index_of_nearest_xy(y_array, x_array, y_point, x_point):
    distance = (y_array-y_point)**2 + (x_array-x_point)**2
    idy,idx = numpy.where(distance==distance.min())
    return idy[0],idx[0]

def do_all(y_array, x_array, points):
    store = []
    for i in xrange(points.shape[1]):
        store.append(find_index_of_nearest_xy(y_array,x_array,points[0,i],points[1,i]))
    return store


# Create some dummy data
y_array = numpy.random.random(10000).reshape(100,100)
x_array = numpy.random.random(10000).reshape(100,100)

points = numpy.random.random(10000).reshape(2,5000)

# Time how long it takes to run
start = time.time()
results = do_all(y_array, x_array, points)
end = time.time()
print 'Completed in: ',end-start

我正在大型数据集上执行此操作，并且真的想加快速度。谁能优化这个？

谢谢。

更新：根据@silvado和@justin（如下）的建议的解决方案

# Shoe-horn existing data for entry into KDTree routines
combined_x_y_arrays = numpy.dstack([y_array.ravel(),x_array.ravel()])[0]
points_list = list(points.transpose())


def do_kdtree(combined_x_y_arrays,points):
    mytree = scipy.spatial.cKDTree(combined_x_y_arrays)
    dist, indexes = mytree.query(points)
    return indexes

start = time.time()
results2 = do_kdtree(combined_x_y_arrays,points_list)
end = time.time()
print 'Completed in: ',end-start

上面的这段代码使我的代码加速了100倍（在100x100矩阵中搜索5000点）。有趣的是，使用scipy.spatial.KDTree（而不是scipy.spatial.cKDTree）给出了与我朴素的解决方案相当的计时，因此使用cKDTree版本绝对值得…

Answer 1

scipy.spatial也有一个kd树实现：scipy.spatial.KDTree。

该方法通常是首先使用点数据来构建kd树。其计算复杂度约为N log
N，其中N是数据点的数量。然后可以进行对数N复杂度的范围查询和最近邻居搜索。这比简单地遍历所有点（复杂度N）要有效得多。

因此，如果您重复了范围查询或最近邻查询，则强烈建议使用kd树。