Python中如何根据numpy数组最大值索引提取另一数组对应元素

解法 #

你已经拿到了axis=0方向的最大值索引max_ind，用numpy内置的take_along_axis方法可以直接提取对应元素，代码如下：

import numpy as np
# 你已有的代码
max_values = array_1.max(axis=0)
max_ind = array_1.argmax(axis=0)
# 提取array_2对应元素
cor_elements = np.take_along_axis(array_2, max_ind[np.newaxis, ...], axis=0)[0]

如果想用高级索引手动实现，也可以用下面的写法：

y, x = np.indices(max_ind.shape)
cor_elements = array_2[max_ind, y, x]

代码说明 #

• take_along_axis方法会按照你传入的索引和指定轴，直接从目标数组中提取对应元素，我们给max_ind加了一个新的轴（和array_2的维度数匹配），提取后去掉多余的轴就能得到形状为(2,4)的结果
• 手动索引的方式中，np.indices生成了后两个维度的坐标网格，和max_ind的第一个维度索引一一对应，即可定位到array_2中所有目标位置的元素

运行上述代码得到的cor_elements就和你给出的预期结果完全一致。

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