如何利用NumPy数组x的最大值索引提取数组y对应元素？

问题原因 #

你用y[idx]得到错误结果，是因为NumPy将idx中的每个值当作行索引来提取整行数据，而非每行对应的列索引。y[[2,0,1]]实际取的是y的第2行、第0行、第1行，所以输出是三行二维数组。

正确解决方案 #

方法1：整数数组索引（最常用）

通过同时指定行索引和列索引，精准定位每个目标元素：

import numpy as np

x = np.array([[2,4,6],
              [9,4,6],
              [6,8,3]])

y = np.array([[88,55,33],
              [43,87,65],
              [98,34,56]])

idx = x.argmax(axis=1)
# 生成与idx长度匹配的行索引（0,1,2）
row_indices = np.arange(x.shape[0])
# 按行+列索引提取元素
result = y[row_indices, idx]

print(result)  # 输出: array([33, 43, 34])

方法2：使用np.take_along_axis（更直观的轴方向提取）

该函数专门用于沿指定轴提取对应索引的元素，无需手动构造行索引：

import numpy as np

x = np.array([[2,4,6],
              [9,4,6],
              [6,8,3]])

y = np.array([[88,55,33],
              [43,87,65],
              [98,34,56]])

idx = x.argmax(axis=1)
# 将idx转为二维数组，匹配y的维度
result = np.take_along_axis(y, idx[:, np.newaxis], axis=1).flatten()

print(result)  # 输出: array([33, 43, 34])

idx[:, np.newaxis]把一维的idx转换为[[2],[0],[1]]，确保和y的二维结构对齐，flatten()则将结果转为一维数组。

方法3：列表推导式（适合新手理解）

通过遍历每行，手动提取对应列的元素：

import numpy as np

x = np.array([[2,4,6],
              [9,4,6],
              [6,8,3]])

y = np.array([[88,55,33],
              [43,87,65],
              [98,34,56]])

idx = x.argmax(axis=1)
result = np.array([y[i, idx[i]] for i in range(len(idx))])

print(result)  # 输出: array([33, 43, 34])

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