如何使用numpy获取二维矩阵每行的前n个最大值及其对应索引

核心思路 #

利用numpy内置的排序索引函数直接按行做降序排序，截取前n位的索引即可，实现逻辑简洁高效，不需要额外拼接偏移量。

完整代码示例 #

首先定义示例矩阵：

import numpy as np 
a = np.array([[10,50,30],[60,20,40],[15,30,90]])

核心实现代码：

# 自定义要取的前n个最大值数量
n = 2
# 按行降序排序后取前n个的列索引
top_n_col_indices = np.argsort(-a, axis=1)[:, :n]

结果验证 #

你可以通过以下代码验证结果正确性：

# 输出每行前n个最大值的列索引
print(top_n_col_indices)
# 输出对应位置的元素值
for row in range(a.shape[0]):
    print(f"第{row+1}行前{n}大的元素：", a[row, top_n_col_indices[row]])

示例输出（n=2时） #

[[1 2]
 [0 2]
 [2 1]]
第1行前2大的元素： [50 30]
第2行前2大的元素： [60 40]
第3行前2大的元素： [90 30]

扩展说明 #

如果你需要分别获取所有前n个元素对应的行索引、列索引数组（和你给出的取最大值示例返回格式一致），可以用以下方式处理：

# 生成对应行索引数组，shape和列索引数组一致
row_indices = np.repeat(np.arange(a.shape[0])[:, None], n, axis=1)
col_indices = top_n_col_indices
# 统一取值验证
print(a[row_indices, col_indices])

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