如何获取二维数组沿第二轴每行的最大值及其对应索引

Python原生实现 #

无需引入第三方依赖，直接遍历每行处理即可，通过enumerate同步获取元素索引和值，配合max函数的key参数按值比较得到目标结果：

# 输入二维数组
input_arr = [
    [1, 5, 3],
    [6, 2, 4],
    [4, 3, 5]
]

result = []
for row in input_arr:
    # 取当前行(索引, 最大值)元组
    max_pair = max(enumerate(row), key=lambda item: item[1])
    result.append(list(max_pair))

print(result)

输出为：[[1, 5], [0, 6], [2, 5]] #

Numpy实现 #

如果是处理大规模数组，推荐使用numpy实现，运算效率远高于原生遍历：

import numpy as np

# 输入转为numpy数组
input_arr = np.array([
    [1, 5, 3],
    [6, 2, 4],
    [4, 3, 5]
])

# 提取每行最大值的索引
max_index = input_arr.argmax(axis=1)
# 提取每行最大值
max_value = input_arr[np.arange(input_arr.shape[0]), max_index]
# 拼接为要求的输出格式
result = np.stack([max_index, max_value], axis=1).tolist()

print(result)

输出和原生实现一致。 #

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