如何在numpy数组上进行nD距离和最近邻计算

如何在numpy数组上进行nD距离和最近邻计算 1.所有距离

• 仅使用

  numpy

天真的方法是：

D = np.sqrt(np.sum((X[:, None, :] - Y[None, :, :])**2, axis = -1))

但是，这会占用大量内存，从而形成一个

(i, j, n)

形中间矩阵，并且非常慢

但是，由于@Divakar（

eucl_dist

package，wiki）的一个技巧，我们可以使用一些代数并

np.einsum

进行如下分解：

(X - Y)**2 = X**2 - 2*X*Y + Y**2

D = np.sqrt(          #  (X - Y) ** 2   np.einsum('ij, ij ->i', X, X)[:, None] +    # = X ** 2        np.einsum('ij, ij ->i', Y, Y)          -    # + Y ** 2        2 * X.dot(Y.T))       # - 2 * X * Y

• Y

  是

  X

与上面类似：

XX = np.einsum('ij, ij ->i', X, X)D = np.sqrt(XX[:, None] + XX - 2 * X.dot(X.T))

请注意，使用这种方法，浮点不精确度会使对角线项与零的偏差非常小。如果需要确保它们为零，则需要显式设置它：

np.einsum('ii->i', D)[:] = 0

• 任何包装

scipy.spatial.distance.cdist

是最直观的内置功能，并且比裸机快得多

numpy

from scipy.spatial.distance import cdistD = cdist(X, Y)

cdist

还可以处理很多距离度量以及用户定义的距离度量（尽管这些度量未优化）。检查上面链接的文档以获取详细信息。

• Y

  是

  X

对于自参考距离，其

scipy.spatial.distance.pdist

工作原理类似于

cdist

，但返回一维压缩距离数组，仅使每个项一次即可节省对称距离矩阵上的空间。您可以使用以下方法将其转换为方阵

squareform

from scipy.spatial.distance import pdist, squareformD_cond = pdist(X)D = squareform(D_cond)

2. K最近的邻居（KNN）

• 仅使用

  numpy

我们可以

np.argpartition

用来获取

k-nearest

索引，并使用索引来获取相应的距离值。因此，

D

当数组保存上面获得的距离值时，我们将有-

if k == 1:    k_i = D.argmin(0)else:    k_i = D.argpartition(k, axis = 0)[:k]k_d = np.take_along_axis(D, k_i, axis = 0)

但是，在缩小数据集之前，不取平方根即可加快速度。

np.sqrt

是计算欧几里得范数最慢的部分，因此我们直到最后都不想这样做。

D_sq = np.einsum('ij, ij ->i', X, X)[:, None] +       np.einsum('ij, ij ->i', Y, Y) - 2 * X.dot(Y.T)if k == 1:    k_i = D_sq.argmin(0)else:    k_i = D_sq.argpartition(k, axis = 0)[:k]k_d = np.sqrt(np.take_along_axis(D_sq, k_i, axis = 0))

现在，

np.argpartition

执行间接分区，并不一定要给我们按排序顺序排列的元素，而只是确保第一个

k

元素是最小的。因此，对于排序后的输出，我们需要使用

argsort

上一步的输出-

sorted_idx = k_d.argsort(axis = 0)k_i_sorted = np.take_along_axis(k_i, sorted_idx, axis = 0)k_d_sorted = np.take_along_axis(k_d, sorted_idx, axis = 0)

如果只需要，则

k_i

根本不需要平方根：

D_sq = np.einsum('ij, ij ->i', X, X)[:, None] +       np.einsum('ij, ij ->i', Y, Y) - 2 * X.dot(Y.T)if k == 1:    k_i = D_sq.argmin(0)else:    k_i = D_sq.argpartition(k, axis = 0)[:k]k_d_sq = np.take_along_axis(D_sq, k_i, axis = 0)sorted_idx = k_d_sq.argsort(axis = 0)k_i_sorted = np.take_along_axis(k_i, sorted_idx, axis = 0)

• X

  是

  Y

在上面的代码中，替换为：

D_sq = np.einsum('ij, ij ->i', X, X)[:, None] +       np.einsum('ij, ij ->i', Y, Y) - 2 * X.dot(Y.T)

与：

XX = np.einsum('ij, ij ->i', X, X)D_sq = XX[:, None] + XX - 2 * X.dot(X.T))

• 任何包装

KD-
Tree
是查找邻居和约束距离的更快方法。请注意，尽管KDTree通常比3d上的强力解决方案要快得多（只要

n

oyu拥有8个以上的点），但如果您具有-
dimensionions，则KDTree仅在拥有多个

2**n

点时才能很好地缩放。有关讨论和更高级的方法，请参见此处

实施KDTree最值得推荐的方法是使用

scipy

的

scipy.spatial.KDTree

或

scipy.spatial.cKDTree

from scipy.spatial.distance import KDTreeX_tree = KDTree(X)k_d, k_i = X_tree.query(Y, k = k)

不幸

scipy

的是，KDTree的实现速度很慢，并且倾向于对较大的数据集进行段错误处理。正如@HansMusgrave指出这里，

pykdtree

增加了性能提升不少，但并不常见的包括作为

scipy

，只能用欧氏距离处理当前（而

KDTree

在

scipy

可以处理任何顺序的Minkowsi
p范）

• X

  是

  Y

改为使用：

k_d, k_i = X_tree.query(X, k = k)

• 任意指标

BallTree具有与KDTree类似的算法属性。我不知道Python中的并行/向量化/快速BallTree，但是使用scipy，我们仍然可以对用户定义的指标进行合理的KNN查询。如果可用，内置指标将更快。

def d(a, b):    return max(np.abs(a-b))tree = sklearn.neighbors.BallTree(X, metric=d)k_d, k_i = tree.query(Y)

如果不是度量标准，则此答案
将是错误的
。BallTree比暴力破解更快的唯一原因是度量标准的属性可以排除某些解决方案。对于真正的任意功能，实际上必须使用蛮力。

d()

3.半径搜索

• 仅使用

  numpy

最简单的方法就是使用布尔索引：

mask = D_sq < r**2r_i, r_j = np.where(mask)r_d = np.sqrt(D_sq[mask])

• 任何包装

与上面类似，您可以使用

scipy.spatial.KDTree.query_ball_point

r_ij = X_tree.query_ball_point(Y, r = r)

要么

scipy.spatial.KDTree.query_ball_tree

Y_tree = KDTree(Y)r_ij = X_tree.query_ball_tree(Y_tree, r = r)

不幸的是

r_ij

最终得到了一个索引数组的列表，这些索引数组很难解开以备后用。

更容易是使用

cKDTree

的

sparse_distance_matrix

，它可以输出

coo_matrix

from scipy.spatial.distance import cKDTreeX_cTree = cKDTree(X)Y_cTree = cKDTree(Y)D_coo = X_cTree.sparse_distance_matrix(Y_cTree, r = r, output_type = `coo_matrix`)r_i = D_coo.rowr_j = D_coo.columnr_d = D_coo.data

这是距离矩阵的一种非常灵活的格式，因为它保留一个实际矩阵（如果转换为

csr

），也可以用于许多矢量化 *** 作。