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


82

我有两个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版本绝对值得...


1
只是一个猜测,但也许kd树会有所帮助。我不知道Python是否有实现。
贾斯汀2012年

无需创建列表和转置“点”。改用数组并遍历索引。
西奥Simier

Answers:


48

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

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

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


1
这看起来很有希望。我将开始阅读有关它的内容,看看是否可以解决问题……
Pete W

1
我仍在测试我的代码,但是早期迹象表明,使用scipy.spatial.cKDTree的速度比我的幼稚方法快100倍。如果明天有更多时间,我将发布最终代码,并且很可能会接受此答案(除非在此之前有更快的方法!)。谢谢你的帮助。
皮特·W

好的,使用scipy.spatial.cKDTree似乎是可行的方法。用我的测试数据进行的测试表明,标准scipy.spatial.KDTree相对于我的幼稚解决方案没有太大/任何改进。
皮特·W

74

这是一个scipy.spatial.KDTree例子

In [1]: from scipy import spatial

In [2]: import numpy as np

In [3]: A = np.random.random((10,2))*100

In [4]: A
Out[4]:
array([[ 68.83402637,  38.07632221],
       [ 76.84704074,  24.9395109 ],
       [ 16.26715795,  98.52763827],
       [ 70.99411985,  67.31740151],
       [ 71.72452181,  24.13516764],
       [ 17.22707611,  20.65425362],
       [ 43.85122458,  21.50624882],
       [ 76.71987125,  44.95031274],
       [ 63.77341073,  78.87417774],
       [  8.45828909,  30.18426696]])

In [5]: pt = [6, 30]  # <-- the point to find

In [6]: A[spatial.KDTree(A).query(pt)[1]] # <-- the nearest point 
Out[6]: array([  8.45828909,  30.18426696])

#how it works!
In [7]: distance,index = spatial.KDTree(A).query(pt)

In [8]: distance # <-- The distances to the nearest neighbors
Out[8]: 2.4651855048258393

In [9]: index # <-- The locations of the neighbors
Out[9]: 9

#then 
In [10]: A[index]
Out[10]: array([  8.45828909,  30.18426696])

5
感谢您提供一个可行的(简单的)示例的完整答案,非常感谢!
johndodo

@lostCrotchet我想是..我还使用了不止一对数据。例如(x,y,z,i)
efirvida

5

如果您可以将数据整理成正确的格式,那么一种快速的方法是使用以下方法scipy.spatial.distance

http://docs.scipy.org/doc/scipy/reference/spatial.distance.html

尤其是pdistcdist提供快速的方法来计算成对距离。


我也称其为按摩,它几乎描述了我们如何处理数据。:D
Lorinc Nyitrai

1
Scipy.spatil.distance是个很棒的工具,但是请注意,如果您有很多距离来计算cKdtree,则它比cdist快得多。
Losbaltica

1
如果我没有误解,使用cdist()或其它numpy的方法示于该答案codereview.stackexchange.com/a/134918/156228
亚历˚F
By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.