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

【问题标题】：Get N maximum values and indices along an axis in a NumPy array在 NumPy 数组中沿轴获取 N 个最大值和索引
【发布时间】：2017-03-27 05:22:50
【问题描述】：

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

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

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

我想获取每个样本的前 M 个项目索引。我也想获得前 M 分数。例如，

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

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

使用 numpy 执行此操作的最佳方法是什么？

【问题讨论】：

@Kasramvd 2D 的链接 dup 目标答案不会按照此问题的需要保持从最高到最低的顺序。所以，我正在重新开放，希望听起来不错。
@Divakar 我的回答提供了一个有序的结果。
@Kasramvd @Kasramvd np.argpartition docs 声明要保持顺序，我们需要为其提供一个整数序列，我认为这将是这些索引的范围，即range(N) 在np.argpartition() 内。你在那里用过这样的东西吗？对不起，如果我错过了！
@Divakar 不，但当然仍然不能证明这不是重复的。另外我认为这种方法与argsort 版本没有任何显着差异。
@Kasramvd 这在很大程度上取决于尺寸。但这里有一个 runtime test ，其中 N 明显小于轴长度。

标签： python numpy matrix

【解决方案1】：

简单的方法是：

获取前 2 个指数

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

获取前 2 个值

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

【讨论】：

【解决方案2】：

如果有人对值和相应的索引都感兴趣而不用调整顺序，下面的简单方法会很有帮助。虽然如果处理大数据可能在计算上很昂贵，因为我们使用 list 来存储 tuples 值，索引。

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

【讨论】：

【解决方案3】：

这是一种使用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 在类似问题上的表现。

【讨论】：

【解决方案4】：

我会使用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]])

【讨论】：