Python中如何将(2,9)的2D NumPy数组重塑为(3,2,3)的3D数组

你可以通过两种高效的NumPy操作实现这个维度转换：

方法1：重塑+转置（推荐，性能最优） #

利用reshape先将原数组的轴1拆分为3个长度为3的分组，再通过transpose调整轴的顺序：

import numpy as np

a = np.array([[ 1,  1,  1, 40, 40, 40, 80, 80, 80],
              [ 0,  0,  1,  0,  1,  1,  0,  1,  1]])

# 先重塑为(2, 3, 3)，再转置轴顺序得到(3, 2, 3)
output = a.reshape(2, 3, 3).transpose(1, 0, 2)
print(output)

• reshape(2, 3, 3)：将原数组的9列拆分为3组，每组3个元素，此时数组形状为(2, 3, 3)，第二个轴对应分组维度
• transpose(1, 0, 2)：将分组维度（原第二个轴）移到第一个位置，最终得到目标形状(3, 2, 3)

方法2：分割+堆叠 #

通过np.split将原数组沿轴1分割为3个子数组，再用np.stack沿新轴堆叠：

import numpy as np

a = np.array([[ 1,  1,  1, 40, 40, 40, 80, 80, 80],
              [ 0,  0,  1,  0,  1,  1,  0,  1,  1]])

# 沿轴1分割为3个(2,3)的数组，再沿轴0堆叠
output = np.stack(np.split(a, 3, axis=1), axis=0)
print(output)

• np.split(a, 3, axis=1)：将原数组的9列平均分割为3个(2,3)的子数组
• np.stack(..., axis=0)：将这3个子数组沿新的轴0堆叠，形成(3, 2, 3)的3D数组

两种方法都能得到你需要的结果，其中方法1的性能更优，因为reshape和transpose都是低开销的数组视图操作，不会额外占用内存。

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