numpy:可变大小数组块中值的有效求和

numpy: efficient summation of values within variably-sized array blocks

问题与中的问题非常相似,我需要按块对矩阵中的元素求和。这里的不同之处在于块可以有不同的大小。例如,给定一个 4×5 矩阵

1 1 | 1 1 1
----|------
1 1 | 1 1 1
1 1 | 1 1 1
1 1 | 1 1 1

行的块大小 1 3 和列的块大小 2 3,结果应该是一个 2×2 矩阵:

2 3
6 9

有没有没有循环的方法?

似乎很适合使用 np.add.reduceat 来基本上沿行求和,然后沿列求和 -

def sum_blocks(a, row_sizes, col_sizes):
    # Sum rows based on row-sizes
    s1 = np.add.reduceat(a,np.r_[0,row_sizes[:-1].cumsum()],axis=0)

    # Sum cols from row-summed output based on col-sizes
    return np.add.reduceat(s1,np.r_[0,col_sizes[:-1].cumsum()],axis=1)

样本运行-

In [45]: np.random.seed(0)
    ...: a = np.random.randint(0,9,(4,5))

In [46]: a
Out[46]: 
array([[5, 0, 3, 3, 7],
       [3, 5, 2, 4, 7],
       [6, 8, 8, 1, 6],
       [7, 7, 8, 1, 5]])

In [47]: row_sizes = np.array([1,3])
    ...: col_sizes = np.array([2,3])

In [48]: sum_blocks(a, row_sizes, col_sizes)
Out[48]: 
array([[ 5, 13],
       [36, 42]])