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

Question

该问题旨在作为规范的重复目标

给定两个阵列X和Y形状(i, n)和(j, n)代表的名单n维坐标，

def test_data(n, i, j, r = 100):
    X = np.random.rand(i, n) * r - r / 2
    Y = np.random.rand(j, n) * r - r / 2
    return X, Y

X, Y = test_data(3, 1000, 1000)

找到最快的方法是：

1. 每个点和每个点之间的D形状距离(i,j)``X``Y
2. 每个点中所有点的最近邻居的索引k_i和距离k_d``k``X``Y
3. 这些指数r_i，r_j以及距离r_d在每一个点的X中距离r每一点的j在Y

鉴于以下几组限制：

• 仅使用 numpy
• 使用任何python包装

包括特殊情况：

• Y 是 X

在所有情况下，距离主要是指欧几里得距离，但是请随意强调允许其他距离计算的方法。

Answer 1

• 仅使用numpy

天真的方法是：

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

但是，这会占用大量内存，从而形成一个(i, j, n)形中间矩阵，并且非常慢

但是，由于@Divakar（eucl_distpackage，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 cdist
D = cdist(X, Y)

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

• Y 是 X

对于自参考距离，其scipy.spatial.distance.pdist工作原理类似于cdist，但返回一维压缩距离数组，仅使每个项一次即可节省对称距离矩阵上的空间。您可以使用以下方法将其转换为方阵squareform

from scipy.spatial.distance import pdist, squareform
D_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上的强力解决方案要快得多（只要noyu拥有8个以上的点），但如果您具有-
dimensionions，则KDTree仅在拥有多个2**n点时才能很好地缩放。有关讨论和更高级的方法，请参见此处

实施KDTree最值得推荐的方法是使用scipy的scipy.spatial.KDTree或scipy.spatial.cKDTree

from scipy.spatial.distance import KDTree
X_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**2
r_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 cKDTree
X_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.row
r_j = D_coo.column
r_d = D_coo.data

这是距离矩阵的一种非常灵活的格式，因为它保留一个实际矩阵（如果转换为csr），也可以用于许多矢量化操作。