获取二维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.        ]])

我需要得到该矩阵 dists 中前N个最大值所对应的索引，因为这些索引将对应于 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个成对距离最大的向量？

- vaer-k

Answer 1

我会先解构、排序，再重新构造。我不是在说这是最好的方法，只是这是我想到的第一种方式。也许在有人发布更简单明了的方案后，我会因羞愧而删除此段文字。 :-)

话虽如此（任意选择前两个值）：

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

- DSM

---

如果您有一个对称的二维数组，这显然是行不通的。 - Moher

Answer 2

作为对DSM非常好的回答的微小改进，如果不考虑N个最大值的顺序，则使用np.argpartition()比使用np.argsort()更有效率。

使用索引i对数组arr进行分区会重新排列元素，使位于索引i处的元素为第i大的元素，而左侧的元素较大，右侧的元素较小。左侧和右侧的分区不一定排序。这样做的优点是它可以在线性时间内运行。

- storm125