Для скорости вы действительно хотите использовать NumPy , который будет значительно быстрее, чем базовый Python, в дополнение к предоставлению более чистого кода для матриц. Из вашего небольшого примера вы можете использовать numpy.cumsum()
дважды по разным осям:
import numpy as np
arr = np.array([[1,1,2],[2,3,4]])
out = arr.cumsum(axis=1).cumsum(axis=0)
print(out)
Дает:
array([[ 1, 2, 4],
[ 3, 7, 13]], dtype=int32)
Примечание: в Windows тип int
по умолчанию - 32-битный, а cumsum()
может молча переполняться при больших матрицах / больших числах, поэтому вам, вероятно, потребуется arr = np.array([[1,1,2],[2,3,4]]).astype(np.int64)
, если в Windows.
Тайминги:
arr = np.arange(10000).reshape(100, 100)
%timeit out = arr.cumsum(axis=1).cumsum(axis=0)
56.3 µs ± 4.96 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Итак, в тысячи раз быстрее, чем требуется.