在numpy数组中按值获取前k个数字的索引？

Question

假设有一个 numpy 数组，如下所示：

import numpy as np
a = array([[1,2,3],[4,3,1]]) 

并且您想返回前 3 个数字。您可以通过以下方式找到最大值：

i,j = np.unravel_index(a.argmax(), a.shape)

但是对于第 2 个和第 3 个最高值的索引，您如何找到它们？我能想到的唯一选择是将存储在最大索引处的值设置为 0，然后再次运行该过程。一定有更好的方法吧？

以下问题的答案是：(1,0)、(0,2)、(1,1)

Answer 1

如果将axis=None参数传递给argsort，它将返回扁平数组的排序索引（按升序）。该numpy.unravel_index函数将扁平数组的索引转换为给定形状数组的索引。

>> a = np.array([[1,2,3],[4,3,1]])
>> np.unravel_index(a.argsort(axis=None), dims=a.shape)
   (array([0, 1, 0, 0, 1, 1], dtype=int64),
    array([0, 2, 1, 2, 1, 0], dtype=int64))

的结果unravel_index是一个数组元组，其中每个数组是沿每个相应轴的索引。要将索引与坐标配对，我们可以使用zip：

>> ix = np.unravel_index(a.argsort(axis=None), dims=a.shape)
>> zip(*ix)[-3:]
   [(0, 2), (1, 1), (1, 0)]

这又是按升序排列的。我们可以使用负步幅来按降序得到它。

>> zip(*ix)[:2:-1]
   [(1, 0), (1, 1), (0, 2)]

这是一行的全部内容，n您想要的“顶部”坐标的数量在哪里。

>> zip(*np.unravel_index(a.argsort(axis=None), dims=a.shape))[::-1][:n]