我有两个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版本绝对值得…
scipy.spatial也有一个kd树实现:scipy.spatial.KDTree。
scipy.spatial
scipy.spatial.KDTree
该方法通常是首先使用点数据来构建kd树。其计算复杂度约为N log N,其中N是数据点的数量。然后可以进行对数N复杂度的范围查询和最近邻居搜索。这比简单地遍历所有点(复杂度N)要有效得多。
因此,如果您重复了范围查询或最近邻查询,则强烈建议使用kd树。