如何从堆叠2D矩阵中提取第二列最大值对应的行？

问题说明 #

现有一个由3个2D矩阵堆叠而成的数组：

[[[-1, 90],
  [-2, 50],
  [-3, 10],
  [-3, 40]],

 [[-4, 99],
  [-5, 40],
  [-6,  5],
  [-3, 50]],

 [[-7,  0],
  [-8,  0],
  [-9, 60],
  [-3, 55]]]

需要返回一个2D矩阵，其中每一行是所有堆叠矩阵中对应行位置里第二列数值最大的那一行。

预期输出：

[[-4, 99],
 [-2, 50],
 [-9, 60],
 [-3, 55]]

直接使用np.max无法得到正确结果，因为它会分别取两列各自的最大值，得到错误输出：

[[-1, 99],
 [-2, 50],
 [-3, 60],
 [-3, 55]]

解决方案 #

利用np.argmax找到每个行位置上第二列最大值所在的堆叠矩阵索引，再根据索引提取对应行即可：

import numpy as np

# 转换为numpy数组
stacked_arr = np.array([
    [[-1, 90], [-2, 50], [-3, 10], [-3, 40]],
    [[-4, 99], [-5, 40], [-6, 5], [-3, 50]],
    [[-7, 0], [-8, 0], [-9, 60], [-3, 55]]
])

# 提取所有堆叠矩阵的第二列，形状为(3, 4)
second_column_vals = stacked_arr[:, :, 1]
# 对每个行位置（axis=0），找到第二列最大值所在的堆叠索引
max_indices = np.argmax(second_column_vals, axis=0)
# 根据索引提取对应行
result = stacked_arr[max_indices, np.arange(stacked_arr.shape[1])]

print(result)

原理说明 #

• second_column_vals提取了所有堆叠矩阵的第二列，每一列对应原数组的一个行位置；
• np.argmax(second_column_vals, axis=0)会针对每个行位置，找出第二列最大值在堆叠维度（第一个维度）的索引；
• 最后用max_indices和行位置索引np.arange(stacked_arr.shape[1])，从原堆叠数组中精准提取到每个行位置的目标行。

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