获取二维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，然后解散。我并不是说这是最好的方法，只是这是我想到的第一种方法，在有人发表更明显的内容后，我可能会羞愧地删除它。:-)

就是说（任意选择前两个值）：

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])