如何使用NumPy获取二维数组每行最大两个值的对应索引

直接用numpy.argsort沿行维度排序取索引即可，代码逻辑简单且执行效率高，完全匹配示例要求。

完整可运行代码 #

import numpy as np

x = np.array([[7, 5, 6],
              [4, 9, 3],
              [1, 6, 7]])

# 沿行方向升序排序取索引，截取最大两个值的索引后调整为从大到小的顺序
row_top2_idx = np.argsort(x, axis=1)[:, -2:][:, ::-1]

# 逐行输出和示例格式一致的结果
for idx_pair in row_top2_idx:
    print(list(idx_pair))

代码说明 #

• 传入axis=1参数指定argsort沿行维度计算，返回值和原数组形状一致，每个位置存的是排序后对应原数组的索引，默认升序排列，值越小的元素对应索引越靠前
• 切片[:, -2:]截取每行最后两位索引，对应每行数值最大的两个元素，此时顺序为「第二大元素索引、最大元素索引」
• 切片[:, ::-1]将两位索引倒序，调整为「最大元素索引、第二大元素索引」，和示例输出顺序完全对齐

运行上述代码的输出结果为：

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

如果不需要严格区分第一大、第二大的顺序，直接去掉[:, ::-1]的倒序步骤即可。如果遇到行内存在重复最大值的场景，argsort会按照元素在原数组中的先后顺序返回索引，符合常规使用预期。

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