在NumPy数组中沿轴获取N个最大值和下标

Question

我认为这对经验丰富的numpy用户来说是一个容易的问题。

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

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

我想要得到每个样本项目的top-M索引。我也想拿到前m的分数。例如,

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

使用numpy最好的方法是什么?

Answer 1

这是使用np.argpartition -

的方法

idx = np.argpartition(a,range(M))[:,:-M-1:-1] # topM_ind
out = a[np.arange(a.shape[0])[:,None],idx]    # topM_score

示例运行-

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

或者，可能更慢，但更短的代码来获得idx将与np.argsort -

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

这是一个post包含一些运行时测试，比较np.argsort和np.argpartition在一个类似的问题。

Answer 2

我会使用argsort():

top2_ind = score_matrix.argsort()[:,::-1][:,:2]

也就是说，生成一个数组，其中包含对score_matrix进行排序的索引:

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

Answer 3

如果有人对值和相应的索引都感兴趣，而不需要调整顺序，那么下面的简单方法将会很有帮助。如果处理大数据，由于使用列表来存储value、index.

的元组，因此在计算上可能会很昂贵。

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

Answer 4

简单的方法是:

获取top-2 索引

np.argsort(-score_matrix)[:, :2]

获取top-2 值

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