python - 在 NumPy 数组中沿轴获取 N 个最大值和索引

Question

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

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

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

我想为每个样本获取项目的前 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

sample 运行-

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 的代码可能更慢，但更短一些

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

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