如何使用Numpy按网格将二维数组分割为多个子数组？

示例数组与需求 #

给定如下6×6的Numpy数组：

import numpy as np

c = np.arange(36).reshape(6, 6)
print(c)

输出：

[[ 0  1  2  3  4  5]
 [ 6  7  8  9 10 11]
 [12 13 14 15 16 17]
 [18 19 20 21 22 23]
 [24 25 26 27 28 29]
 [30 31 32 33 34 35]]

我们需要将其按3×3网格分割为9个2×2的子数组，最终得到包含9个二维数组的结构：

[[[0, 1], [6, 7]], 
 [[2, 3], [8, 9]],
 [[4, 5], [10, 11]],
 [[12, 13], [18, 19]],
 [[14, 15], [20, 21]],
 [[16, 17], [22, 23]],
 [[24, 25], [30, 31]],
 [[26, 27], [32, 33]],
 [[28, 29], [34, 35]]]

通用实现方法 #

对于任意可以被(N,M)网格整除的二维数组（即原数组行数能被N整除，列数能被M整除），可以通过Numpy的维度变换高效实现分割：

步骤1：定义参数

• N：网格的行数（示例中为3）
• M：网格的列数（示例中为3）
• sub_row：每个子数组的行数，计算公式：sub_row = 原数组行数 // N
• sub_col：每个子数组的列数，计算公式：sub_col = 原数组列数 // M

步骤2：代码实现

# 定义网格参数
N, M = 3, 3
# 获取子数组的尺寸
sub_row, sub_col = c.shape[0] // N, c.shape[1] // M
# 执行分割
result = c.reshape(N, sub_row, M, sub_col).transpose(0, 2, 1, 3).reshape(-1, sub_row, sub_col)

# 打印结果
print(result)

逻辑解释

1. 维度拆分：通过reshape(N, sub_row, M, sub_col)将原数组拆分为四维结构：第一维度对应网格的行分组，第二维度是子数组的行，第三维度对应网格的列分组，第四维度是子数组的列。
2. 调整维度顺序：transpose(0,2,1,3)将维度顺序调整为(网格行, 网格列, 子数组行, 子数组列)，让每个网格位置的子数组数据连续排列。
3. 合并网格维度：最后用reshape(-1, sub_row, sub_col)将前两个网格维度合并，得到包含N*M个子数组的三维数组，每个子数组保持(sub_row, sub_col)的形状。

注意事项

• 若原数组的行数/列数无法被N/M整除，需先对数组进行裁剪或填充操作，确保尺寸匹配后再执行分割。
• 该方法基于Numpy的视图操作，不会复制原始数据，运行效率极高。

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