如何用更Pythonic的方式获取多个同尺寸NumPy数组的逐元素极值

要实现这个需求，完全不需要嵌套循环，用NumPy的数组堆叠和索引操作就能高效完成，步骤如下：

• 把所有输入数组沿新的轴（比如第三轴）堆叠，让每个位置的候选元素集中在同一维度
• 计算每个位置上元素绝对值最大的对应索引
• 根据索引从堆叠后的数组中提取目标元素

完整代码示例 #

import numpy as np

a = np.array([[-1, 2], [3, 4]])
b = np.array([[-3, 1], [5, -2]])
c = np.array([[2, -1], [1, -5]])

# 1. 沿第三轴堆叠数组，得到(m, n, 3)形状的三维数组
stacked = np.stack([a, b, c], axis=2)
# 2. 找出每个位置绝对值最大的元素的索引
abs_max_indices = np.argmax(np.abs(stacked), axis=2)
# 3. 根据索引提取对应元素，还原为原二维形状
out = np.take_along_axis(stacked, abs_max_indices[..., np.newaxis], axis=2).squeeze(axis=2)

print(out)
# 输出：[[-3  2]
#        [ 5 -5]]

为什么np.maximum.reduce/np.minimum.reduce无法满足需求？ #

这两个方法只能分别取出对应位置的最大值或最小值，但我们要的是绝对值最大的元素——它可能是正数里的最大值，也可能是负数里的最小值（比如示例中的-3、-5）。单纯取最大或最小都无法覆盖所有情况，因此必须先通过绝对值定位目标元素的位置，再提取原元素。

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