在2D numpy ndarray或numpy矩阵中获取前N个值的索引

Question

我有一组N维向量.

data = np.array([[5, 6, 1], [2, 0, 8], [4, 9, 3]])

In [1]: data
Out[1]:
array([[5, 6, 1],
       [2, 0, 8],
       [4, 9, 3]])

我正在使用sklearn pairwise_distances函数来计算距离值矩阵.注意,该矩阵关于对角线是对称的.

dists = pairwise_distances(data)

In [2]: dists
Out[2]:
array([[  0.        ,   9.69535971,   3.74165739],
       [  9.69535971,   0.        ,  10.48808848],
       [  3.74165739,  10.48808848,   0.        ]])

我需要与该矩阵中的前N个值dists对应的索引,因为这些索引将对应于data表示它们之间距离最大的向量的成对索引.

我试着np.argmax(np.max(distances, axis=1))去获取每一行中最大值的索引,并np.argmax(np.max(distances, axis=0))获得每列中最大值的索引,但请注意:

In [3]: np.argmax(np.max(dists, axis=1))
Out[3]: 1

In [4]: np.argmax(np.max(dists, axis=0))
Out[4]: 1

和:

In [5]: dists[1, 1]
Out[5]: 0.0

因为矩阵关于对角线是对称的,并且因为argmax返回它找到的具有最大值的第一个索引,所以我最终得到行中的对角线中的单元格与存储最大值的列匹配,而不是行和最高值的列本身.

在这一点上,我确信我可以编写更多代码来找到我正在寻找的值,但肯定有一种更简单的方法可以做我想做的事情.所以我有两个或多或少等同的问题:

如何在矩阵中找到与前N个值对应的索引,或者,如何从向量数组中找到具有前N个成对距离的向量？

Answer 1

我拉扯,argsort,然后解开.我并不是说这是最好的方式,只是这是我发生的第一种方式,而且在有人发布更明显的内容之后,我可能会羞愧地删除它.:-)

那说(选择前2个值,任意):

In [73]: dists = sklearn.metrics.pairwise_distances(data)

In [74]: dists[np.tril_indices_from(dists, -1)] = 0

In [75]: dists
Out[75]: 
array([[  0.        ,   9.69535971,   3.74165739],
       [  0.        ,   0.        ,  10.48808848],
       [  0.        ,   0.        ,   0.        ]])

In [76]: ii = np.unravel_index(np.argsort(dists.ravel())[-2:], dists.shape)

In [77]: ii
Out[77]: (array([0, 1]), array([1, 2]))

In [78]: dists[ii]
Out[78]: array([  9.69535971,  10.48808848])