Python中的矢量化平均K-最近邻距离

时间:2014-05-21 21:22:16

标签: python numpy vectorization

这是R n 中点的K-最近邻算法,应该计算每个点与其k-最近邻居的平均距离。问题在于,虽然它是矢量化的,但在我重复自己的意义上,它的效率很低。如果有人能帮我改进这段代码,我会很高兴:

import numpy as np
from scipy.spatial.distance import pdist
from scipy.spatial.distance import squareform

def nn_args_R_n_squared(points):
    """Calculate pairwise distances of points and return the matrix together with matrix of indices of the first matrix sorted"""
    dist_mat=squareform(pdist(points,'sqeuclidean'))
    return dist_mat,np.argsort(dist_mat,axis=1)
def knn_avg_dist(X,k):
    """Calculates for points in rows of X, the average distance of each, to their k-nearest      neighbours"""
    X_dist_mat,X_sorted_arg=nn_args_R_n_squared(X)
    X_matrices=(X[X_sorted_arg[:,1:k+1]]-X[...,None,...]).astype(np.float64)
    return np.mean(np.linalg.norm(X_matrices,axis=2)**2,axis=1)
X=np.random.randn(30).reshape((10,3))
print X
print knn_avg_dist(X,3)

输出:

[[-1.87979713  0.02832699  0.18654558]
 [ 0.95626677  0.4415187  -0.90220505]
 [ 0.86210012 -0.88348927  0.32462922]
 [ 0.42857316  1.66556448 -0.31829065]
 [ 0.26475478 -1.6807253  -1.37694585]
 [-0.08882175 -0.61925033 -1.77264525]
 [-0.24085553  0.64426394 -0.01973027]
 [-0.86926425  0.93439913 -0.31657442]
 [-0.30987468  0.02925649 -1.38556347]
 [-0.41801804  1.40210993 -1.04450895]]
[ 3.37983833  2.1257945   3.60884158  1.67051682  2.85013297  1.66756279
  1.2678029   1.20491026  1.54623574  1.30722388]

正如您所看到的,我计算了两次距离,但我无法想出一种从X_dist_mat读取相同信息的方法,因为我必须同时读取每行中的多个元素

1 个答案:

答案 0 :(得分:3)

使用scipy.spatial.cKDTree

>>> data = np.random.rand(1000, 3)
>>> import scipy.spatial

>>> kdt = scipy.spatial.cKDTree(data) 
>>> k = 5 # number of nearest neighbors 
>>> dists, neighs = kdt.query(data, k+1)
>>> avg_dists = np.mean(dists[:, 1:], axis=1)