在多维数组中返回N个最大值的索引（可以找到一维数组的解决方案，但不适用于多维数组）。

Question

我有一个numpy数组X，我想返回另一个数组Y，其条目是X中n个最大值的索引，即假设我有：

a =np.array[[1, 3, 5], [4, 5 ,6], [9, 1, 7]]

如果我想要前5个"maxs"的索引-这里的最大值为9、7、6、5、5，它们的索引分别是：

b=np.array[[2, 0], [2 2], [ 2 1],  [1 1], [0 , 2])

我能找到一些解决方案，使这个方法对于一个一维数组起作用，比如：

c=np.array[1, 2, 3, 4, 5, 6]:

def f(a,N): 
    return np.argsort(a)[::-1][:N] 

但是我们还没有能够在多个维度上生成可用的东西。谢谢！

- user7351362

---

谢谢，今天一直很艰难哈哈。 - user7351362

Answer 1

方法 #1

获取其展平版本的argsort索引，并选择最后N个索引。然后，获取相应的行和列索引 -

N = 5
idx = np.argsort(a.ravel())[-N:][::-1] #single slicing: `[:N-2:-1]`
topN_val = a.ravel()[idx]
row_col = np.c_[np.unravel_index(idx, a.shape)]

样例运行 -

# Input array
In [39]: a = np.array([[1,3,5],[4,5,6],[9,1,7]])

In [40]: N = 5
    ...: idx = np.argsort(a.ravel())[-N:][::-1]
    ...: topN_val = a.ravel()[idx]
    ...: row_col = np.c_[np.unravel_index(idx, a.shape)]
    ...: 

In [41]: topN_val
Out[41]: array([9, 7, 6, 5, 5])

In [42]: row_col
Out[42]: 
array([[2, 0],
       [2, 2],
       [1, 2],
       [1, 1],
       [0, 2]])

---

方法 #2

为了提高性能，我们可以使用np.argpartition来获取前N个索引，而不必保持排序顺序，具体如下-

idx0 = np.argpartition(a.ravel(), -N)[-N:]

为了获得排序后的顺序，我们需要进行一轮argsort操作 -

idx = idx0[a.ravel()[idx0].argsort()][::-1]

- Divakar