【问题标题】：numpy ND array: take index'ed elements along a certain axis (ND choose)numpy ND 数组：沿某个轴获取索引元素（ND 选择）
【发布时间】：2016-08-25 12:46:34
【问题描述】：

我沿某个轴获得索引。例如像这样使用 2D 和 axis=-1 ：

>>> axis = -1
>>> a = rand(5, 3) - 0.5;  a
array([[ 0.49970414, -0.14251437,  0.2881351 ],
       [ 0.3280437 ,  0.33766112,  0.4263927 ],
       [ 0.37377502,  0.05392274, -0.4647834 ],
       [-0.09461463, -0.25347861, -0.29381079],
       [-0.09642799,  0.15729681,  0.06048399]])
>>> axisinds = a.__abs__().argmax(axis);  axisinds
array([0, 2, 2, 2, 1])

现在如何通过沿该轴获取索引元素将数组减少一维？

对于 2D 和 axis=-1，可以这样做（为了获得示例数组每行的绝对最大值）：

>>> a[arange(len(axisinds)), axisinds]
array([ 0.49970414,  0.4263927 , -0.4647834 , -0.29381079,  0.15729681])

但这非常特殊，仅限于 1 或 0 个结果维度。 ndim 和 axis 怎么样？

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

现在我自己找到了一个简单的解决方案：

def choose_axis(inds, a, axis=-1):
    return np.choose(inds, np.rollaxis(a, axis))

>>> choose_axis(axisinds, a, -1)
array([ 0.49970414,  0.4263927 , -0.4647834 , -0.29381079,  0.15729681])

编辑： 然而，由于np.choose 的（未记录的）限制，这种方法在轴方向（32 位？）上被限制为最多 31 个元素。在许多情况下，这是可以的。

这是一个

无限方法：

def choose_axis(inds, a, axis=-1):
    # handles any number & size of dimensions, and any axis
    if (axis + 1) % a.ndim:            # move axis to last dim
        a = np.moveaxis(a, axis, -1)   # = np.rollaxis(a, axis, a.ndim)
    shape = a.shape
    a = a.reshape(-1, shape[-1])   # 2D
    a = a[np.arange(inds.size), inds.ravel()]  # effective reduction
    return a.reshape(shape[:-1])

因此，ND 绝对最小值示例可以这样完成：

def absminND(a, axis=-1):
    inds = a.__abs__().argmin(axis)
    if axis is None:
        return a.ravel()[inds]
    return choose_axis(inds, a)

【讨论】：