Numpy二维数组分块求和算法

早上有同事问了一个问题：如何将m行n列二维numpy数组切成若干个u行v列的小块，并求每一小块的元素和，其中m是u的整倍数，n是v的整倍数——简单说就是二维数组分块求和问题。

乍一听，似乎不难，细想却发现这个问题并不简单。梳理了一遍自己熟悉的知识库，numpy没有提供直接解决问题的函数，scipy的二维卷积有点似是而非。看来只能自己构思了，于是我尝试写出了下面这样一段代码。

import numpy
m, n = 4, 6
u, v = 2, 2
arr = np.arange(m*n).reshape(m, n)
arr
array([[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23]])

先构造一个4行6列的二维数组，再将其切分成6个2行2列的小块，并求和。

np.stack([np.sum(np.split(a, n//v, axis=1), axis=(1,2)) for a in np.split(arr, m//u)])
array([[14, 22, 30],
       [62, 70, 78]])

结果看起来是正确的，不过列表推导式本质上还是for循环，这个代码跑起来速度是个大问题。苦思冥想之时灵光闪现，得到这样一个方案：

np.sum(arr.reshape(m//u, u, n//v, v), axis=(1,3))
array([[14, 22, 30],
       [62, 70, 78]])

这应该是最简洁、最高效的方案了。