python - NumPy 沿着不相交的索引求和

Question

我有一个应用程序，我需要对 3D NumPy 数组中的任意索引组求和。内置的 NumPy 数组求和例程对 ndarray 的一个维度上的所有索引求和。相反，我需要对数组中一个维度的索引范围求和并返回一个新数组。

例如，假设我有一个形状为 (70,25,3) 的 ndarray。我希望总结某些索引范围内的第一个维度并返回一个新的 3D 数组。考虑 0:25, 25:50 和 50:75 的总和，它将返回形状为 (3,25,3) 的数组.

是否有一种简单的方法可以沿 NumPy 数组的一个维度进行“不相交求和”以产生此结果？

Answer 1

您可以使用 np.add.reduceat 作为解决此问题的通用方法。即使范围的长度不完全相同，这也有效。

要沿轴 0 对切片 0:25、25:50 和 50:75 求和，传入索引 [ 0, 25, 50]:

np.add.reduceat(a, [0, 25, 50], axis=0)

此方法也可用于对非连续范围求和。例如，要对切片 0:25、37:47 和 51:75 求和，写:

np.add.reduceat(a, [0,25, 37,47, 51], axis=0)[::2]

另一种对相同 长度范围求和的方法是 reshape 数组，然后沿轴求和。与上面的第一个示例等效的是:

a.reshape(3, a.shape[0]//3, a.shape[1], a.shape[2]).sum(axis=1)