Python 沿NumPy数组中的轴获取N个最大值和索引

我认为这对于经验丰富的numpy用户来说是一个简单的问题

我有一个分数矩阵。原始索引对应于样本，列索引对应于项目。比如说,

score_matrix = 
  [[ 1. ,  0.3,  0.4],
   [ 0.2,  0.6,  0.8],
   [ 0.1,  0.3,  0.5]]

top2_ind = 
  [[0, 2],
   [2, 1],
   [2, 1]]

top2_score = 
  [[1. , 0.4],
   [0,8, 0.6],
   [0.5, 0.3]]

我想为每个样本获取项目的top-M索引。我还想得到top-M分数。比如说,

score_matrix = 
  [[ 1. ,  0.3,  0.4],
   [ 0.2,  0.6,  0.8],
   [ 0.1,  0.3,  0.5]]

top2_ind = 
  [[0, 2],
   [2, 1],
   [2, 1]]

top2_score = 
  [[1. , 0.4],
   [0,8, 0.6],
   [0.5, 0.3]]

使用numpy实现这一点的最佳方法是什么？

下面是一种使用-

样本运行-

In [343]: a
Out[343]: 
array([[ 1. ,  0.3,  0.4],
       [ 0.2,  0.6,  0.8],
       [ 0.1,  0.3,  0.5]])

In [344]: M = 2

In [345]: idx = np.argpartition(a,range(M))[:,:-M-1:-1]

In [346]: idx
Out[346]: 
array([[0, 2],
       [2, 1],
       [2, 1]])

In [347]: a[np.arange(a.shape[0])[:,None],idx]
Out[347]: 
array([[ 1. ,  0.4],
       [ 0.8,  0.6],
       [ 0.5,  0.3]])

或者，使用

np.argsort

-

idx = a.argsort(1)[:,:-M-1:-1]

下面是一个包含一些运行时测试的测试，它比较了

np.argsort

和

np.argpartition

在类似问题上的差异。

我会使用

argsort（）

：

也就是说，生成一个包含索引的数组，该索引将排序

得分矩阵

：

array([[1, 2, 0],
       [0, 1, 2],
       [0, 1, 2]])

然后使用

：：-1将列反转，然后使用
：2将前两列反转：

array([[0, 2],
       [2, 1],
       [2, 1]])

然后类似，但使用常规的
np.sort（）
获取值：

top2_score = np.sort(score_matrix)[:,::-1][:,:2]

它遵循与上述相同的机制，为您提供：

array([[ 1. ,  0.4],
       [ 0.8,  0.6],
       [ 0.5,  0.3]])

如果有人对数值和相应的指数都感兴趣，而没有对顺序进行调整，下面的简单方法将很有帮助。但是，如果处理大数据，计算成本可能会很高，因为我们使用的是一个列表来存储值、索引的元组

import numpy as np
values = np.array([0.01,0.6, 0.4, 0.0, 0.1,0.7, 0.12]) # a simple array
values_indices = [] # define an empty list to store values and indices
while values.shape[0]>1:
    values_indices.append((values.max(), values.argmax()))
    # remove the maximum value from the array:
    values = np.delete(values, values.argmax())

最终输出为元组列表：

values_indices
[(0.7, 5), (0.6, 1), (0.4, 1), (0.12, 3), (0.1, 2), (0.01, 0)]

@Kasramvd 2D的链接dup目标答案没有按照此问题的需要保持从最高到最低的顺序。所以，我正在重新打开，希望听起来没问题。@Divakar我的答案提供了一个有序的结果。@Kasramvd很好地说明，为了保持顺序，我们需要提供一个整数序列，我认为这是这些索引的范围，即
range（N）
inside
np.argpartition（）
。你在那里用过这样的东西吗？对不起，我错过了@Divakar没有，但它仍然不能证明这不是一个复制品。另外，我不认为这种方法与
argsort
version有任何显著区别，这在很大程度上取决于大小。但这里是a，其中
N
明显小于轴长度。