numpy: эффективное суммирование значений в блоках массива переменного размера

Question

Проблема очень похожа на проблему в Как оценить сумму значений в блоках массива , где мне нужно суммировать элементы в матрице по блокам. Что здесь отличается, так это то, что блоки могут иметь разные размеры. Например, учитывая матрицу 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

Есть ли способ сделать это без циклов?

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)

Пробный прогон -

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