python - numpy:可变大小数组 block 中值的有效求和

Question

问题与How to evaluate the sum of values within array blocks中的问题非常相似，我需要按 block 对矩阵中的元素求和。这里的不同之处在于 block 可以有不同的大小。例如，给定一个 4×5 矩阵

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

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

2 3
6 9

有没有一种不用循环的方法？

Answer 1

似乎很适合使用 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)

sample 运行-

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]])