如何不使用循环，通过Numpy内置方法生成n*n子矩阵和的新矩阵

我有一个初始数组a，现在通过双层循环生成新矩阵b，其中每个元素对应原数组中所有不重复的n×n子矩阵的和（比如0+1+5+6=12、1+2+6+7=16这类）。能不能不用循环，只靠Numpy内置方法实现相同效果？我试过np.reshape，但处理不了相交的子矩阵。

原循环实现代码 #

import numpy as np
a = np.array([[ 0,  1,  2,  3, 4], [5,  6,  7,  8, 9], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19], [15, 16, 17, 18, 20]])
n = 2

b = np.empty([0, np.size(a[0])-n+1], int)
for row in range(0, np.size(a[0])-n+1):
    temp_list = []
    for col in range(0, np.size(a[0])-n+1):
        temp_list.append(np.sum(a[row:row+n, col:col+n]))
    b = np.vstack([b, [np.array(temp_list)]])
    
print('Input (a):')
print(a)
print('Output (b):')
print(b)

原运行结果 #

Input (a):
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [15 16 17 18 20]]
Output (b):
[[12 16 20 24]
 [32 36 40 44]
 [52 56 60 64]
 [62 66 70 75]]

无循环的Numpy实现方案 #

可以用numpy.lib.stride_tricks.sliding_window_view直接生成所有n×n的滑动子矩阵，再对每个子矩阵求和，完全不需要循环：

import numpy as np

a = np.array([[ 0,  1,  2,  3, 4], [5,  6,  7,  8, 9], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19], [15, 16, 17, 18, 20]])
n = 2

# 生成所有n×n的滑动窗口
windows = np.lib.stride_tricks.sliding_window_view(a, window_shape=(n, n))
# 对每个窗口的最后两个维度求和，得到结果矩阵b
b = windows.sum(axis=(2, 3))

print('Input (a):')
print(a)
print('Output (b):')
print(b)

验证结果 #

运行上述代码，输出和原循环实现完全一致：

Input (a):
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [15 16 17 18 20]]
Output (b):
[[12 16 20 24]
 [32 36 40 44]
 [52 56 60 64]
 [62 66 70 75]]

说明 #

• sliding_window_view会在原数组的基础上生成滑动窗口视图（不复制数据，效率很高），窗口形状指定为(n,n)后，得到的数组形状是(M, N, n, n)，其中M = a.shape[0]-n+1，N = a.shape[1]-n+1，对应原循环中行列的循环次数。
• 最后通过sum(axis=(2,3))对每个窗口的n×n元素求和，直接得到目标矩阵b。

内容的提问来源于stack exchange，提问作者yan