numpy:沿除最后一个轴以外的所有轴求和

时间:2018-08-06 15:35:55

标签: python numpy

如果我有一个任意形状的ndarray,并且我想沿除最后一个轴之外的所有轴计算和,例如,可以通过这样做来实现

all_but_last = tuple(range(arr.ndim - 1))
sum = arr.sum(axis=all_but_last)

现在,tuple(range(arr.ndim - 1))并不完全直观。有没有更优雅/麻木的方式来做到这一点?

此外,如果要对形状不同的多个数组执行此操作,则必须为每个数组计算一个单独的维元组。有没有一种更规范的说法“不管尺寸是多少,只给我一个轴就行”?

2 个答案:

答案 0 :(得分:4)

您可以调整数组的形状,以使除最后一个轴外的所有轴均变平(例如,形状(k, l, m, n)变为(k*l*m, n)),然后在第一个轴上求和。

例如,这是您的计算:

In [170]: arr.shape
Out[170]: (2, 3, 4)

In [171]: arr.sum(axis=tuple(range(arr.ndim - 1)))
Out[171]: array([2.85994792, 2.8922732 , 2.29051163, 2.77275709])

这是替代方案:

In [172]: arr.reshape(-1, arr.shape[-1]).sum(axis=0)
Out[172]: array([2.85994792, 2.8922732 , 2.29051163, 2.77275709])

答案 1 :(得分:3)

您可以使用np.apply_over_axes对多个轴求和。

np.apply_over_axes(np.sum, arr, [0,2]) #sum over axes 0 and 2

np.apply_over_axes(np.sum, arr, range(arr.ndim - 1)) #sum over all but last axis