python - 获得沿轴绝对值的最大值

Question

我有几个形状相同的ndarrays，我想得到一个数组(形状相同)，每个元素的绝对值都是最大值。所以我决定堆叠所有数组，然后沿新的堆叠轴选取值。但是如何做到这一点呢？

例子

假设我们有两个一维数组，每个数组有 4 个元素，所以我的堆叠数组看起来像

>>> stack
array([[ 4,  1,  2,  3],
       [ 0, -5,  6,  7]])

如果我只对我能做到的最大程度感兴趣

>>> numpy.amax(stack, axis=0)
array([4, 1, 6, 7])

但是我也需要考虑负值，所以我打算

>>> ind = numpy.argmax(numpy.absolute(stack), axis=0)
>>> ind
array([0, 1, 1, 1])

现在我有了所需的索引，但如何将其应用于堆叠数组？如果我只是通过 ind 索引 stack，numpy 正在做一些我不需要的广播内容:

>>> stack[ind]
array([[ 4,  1,  2,  3],
       [ 0, -5,  6,  7],
       [ 0, -5,  6,  7],
       [ 0, -5,  6,  7]])

我想要得到的是array([4, -5, 6, 7])

或者从稍微不同的角度来问:如何根据 numpy.argmax(stack,轴=0)?

Answer 1

堆叠操作效率低下。我们可以简单地使用 np.where 来根据绝对值比较进行选择 -

In [198]: a
Out[198]: array([4, 1, 2, 3])

In [199]: b
Out[199]: array([ 0, -5,  6,  7])

In [200]: np.where(np.abs(a) > np.abs(b), a, b)
Out[200]: array([ 4, -5,  6,  7])

这适用于通用的 n-dim 数组，无需任何修改。