Python 3.x 返回N个最大值'；多维数组中的索引（可以找到一维的解决方案，但不能找到多维的解决方案）

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

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

然后说，如果我想要前5个“maxs”的索引，这里9，7，6，5，5是maxs，它们的索引是：

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] 

但我们无法在多个维度上产生效果。谢谢

方法#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]

---

谢谢，这是一个艰难的一天到目前为止哈哈