如何获取数组第五列最大值的对应行索引？

给定如下数组A：

A = [[0.0, 1.0, 3.0, -1.0, -1008.0],
[0.0, 1.0, 3.0, -1.0, -1008.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, -2.0, -1488.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, -4.0, -1808.0],
[0.0, 1.0, 3.0, -2.0, -1488.0],
[0.0, 1.0, 3.0, -3.0, -1648.0],
[0.0, 1.0, 3.0, -3.0, -1648.0],
[0.0, 1.0, 3.0, -3.0, -1648.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, 3.0, -1648.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, 2.0, -1488.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, -2.0, -1488.0],
[0.0, 1.0, 3.0, 2.0, -1488.0],
[0.0, 1.0, 3.0, 3.0, -1648.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, 2.0, -1488.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, -3.0, -1648.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, -1.0, -1328.0],
[0.0, 1.0, 3.0, -2.0, -1488.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, -2.0, -1488.0],
[0.0, 1.0, 3.0, -3.0, -1648.0],
[0.0, 1.0, 3.0, -2.0, -1488.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, 1.0, -1328.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, 0.0, -1168.0],
[0.0, 1.0, 3.0, 0.0, -1168.0]]

需求是找到**第5列（索引为4）**中最大值的所在行索引，期望结果为m = [0, 1]（对应两个值为-1008.0的行）。

原代码错误原因 #

原代码m = np.argwhere(A == np.amax(A))的问题在于：

• np.amax(A)会计算整个二维数组的最大值，而非指定第4列的最大值，导致匹配结果不符合预期。

正确解决方案 #

步骤如下：

1. 将列表形式的A转换为numpy数组（方便索引操作）
2. 提取第4列的数据
3. 计算该列的最大值
4. 找出该列中等于最大值的所有行索引

具体代码：

import numpy as np

# 转换为numpy数组
arr = np.array(A)
# 提取第4列（索引4）
col_4 = arr[:, 4]
# 计算该列的最大值
max_val = np.amax(col_4)
# 获取所有等于最大值的行索引，提取行号并转为列表
m = np.argwhere(col_4 == max_val).flatten().tolist()

print(m)  # 输出: [0, 1]

简化写法 #

也可以合并为一行代码：

m = np.argwhere(np.array(A)[:,4] == np.amax(np.array(A)[:,4])).flatten().tolist()

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