如何将扁平化Numpy数组重塑为指定数量的等维度矩阵？

核心思路 #

要实现按间隔抽取元素分配到多个等维度矩阵，关键是先按目标矩阵的数量对原数组元素分组，再将每组重塑为目标矩阵的形状。

针对示例的实现步骤 #

原输入数组：

import numpy as np
arr = np.array([[2, 2, 1, 0, 4, 2, 3, 3]])

1. 先将二维数组摊平为一维：

flat_arr = arr.ravel()  # 得到 array([2, 2, 1, 0, 4, 2, 3, 3])

2. 假设目标矩阵数量为k=2，每个矩阵形状为(2,2)：
   • 按间隔k抽取元素，分成k组
   • 每组元素数为总元素数除以k，这里8/2=4，刚好能重塑为2x2

k = 2
matrix_shape = (2, 2)
# 按索引分组：第0组取0,2,4,6；第1组取1,3,5,7
groups = [flat_arr[i::k] for i in range(k)]
# 将每组重塑为目标矩阵，再组合成三维数组
result = np.array([group.reshape(matrix_shape) for group in groups])

运行后得到预期结果：

array([[[2, 1],
        [4, 3]],

       [[2, 0],
        [2, 3]]])

扩展到任意数量等维度矩阵 #

只要保证原数组总元素数 = 目标矩阵数量 × 单个矩阵元素数，这个方法通用：

def reshape_by_interval(arr, k, matrix_shape):
    flat_arr = arr.ravel()
    # 验证元素数量是否匹配
    if flat_arr.size != k * np.prod(matrix_shape):
        raise ValueError("元素总数必须等于矩阵数量乘以单个矩阵的元素数")
    groups = [flat_arr[i::k] for i in range(k)]
    return np.array([g.reshape(matrix_shape) for g in groups])

比如要生成3个3x3的矩阵（总元素数需为27），直接调用函数即可：

# 示例：3个3x3矩阵
test_arr = np.arange(27).reshape(1,27)
result = reshape_by_interval(test_arr, 3, (3,3))

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