python - numpy 数组的每一行的特定列组的总和

Question

我有以下 NumPy 数组:

>>> a = np.arange(21).reshape(3, 7)
>>> a
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12, 13],
       [14, 15, 16, 17, 18, 19, 20]])

我想通过用它们的总和替换一些列来创建一个新数组。我需要总结的列将始终组合在一起。
例如，保持前两列不变。将索引为 2 的列替换为索引为 2、3 和 4 的列的总和。将索引为 3 的列替换为索引为 5 和 6 的列的总和。
所需输出:

array([[ 0,  1,  9, 11],
       [ 7,  8, 30, 25],
       [14, 15, 51, 39]])

我试过的:

b = np.concatenate([
    a[:, :2], 
    a[:, 2:5].sum(axis=1, keepdims=True), 
    a[:, 5:].sum(axis=1, keepdims=True)
], axis=1)

这提供了所需的输出，但我想知道是否有更好/简洁的方法来做到这一点。

Answer 1

定义要求和的区间的区间，然后使用 np.add.reduceat 沿axis=1 (对于每一行)-

In [37]: a
Out[37]: 
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12, 13],
       [14, 15, 16, 17, 18, 19, 20]])

In [38]: bins = [0,1,2,5]

In [39]: np.add.reduceat(a,bins,axis=1)
Out[39]: 
array([[ 0,  1,  9, 11],
       [ 7,  8, 30, 25],
       [14, 15, 51, 39]])