如何用Python获取NumPy矩阵每行中N个最小值的索引？

给定如下NumPy矩阵：

import numpy as np
m = np.array([[1, 2, 3, 4],
              [10, 5, 3, 4],
              [12, 8, 1, 2],
              [7, 0, 2, 4]])

需要提取每行中指定数量（例如N=2）最小值的索引，且索引按对应值的升序排列，期望输出：

[[0, 1],
 [2, 3],
 [3, 2],
 [1, 2]]

解决方案 #

直接使用np.argsort()函数即可实现需求，代码如下：

import numpy as np

m = np.array([[1, 2, 3, 4],
              [10, 5, 3, 4],
              [12, 8, 1, 2],
              [7, 0, 2, 4]])
N = 2

# 按行排序获取索引，取每行前N个元素
min_indices = np.argsort(m, axis=1)[:, :N]
print(min_indices)

代码说明 #

• np.argsort(m, axis=1)：指定按行（axis=1）对矩阵元素进行升序排序，返回的矩阵中每个元素是原矩阵对应位置的索引，即每行的索引按对应值从小到大排列
• [:, :N]：对排序后的索引矩阵进行切片，保留每行的前N个元素，也就是每行最小的N个值对应的索引，完全符合需求的升序排列要求

内容的提问来源于stack exchange，提问作者Fredrik