如何对Numpy 4D数组按批次横向拼接子矩阵并调整形状？

要实现将形状为(a1, a2, a3, a4)的4D数组（a1个批次，每个批次含a2个(a3,a4)的2D矩阵）按批次沿最后一维横向拼接，最终得到(a1, a3, a4*a2)的数组，有两种高效方法：

方法一：使用np.concatenate直接拼接 #

通过提取每个批次内的所有矩阵，沿最后一维（列方向）拼接：

import numpy as np

# 示例数组
a1, a2, a3, a4 = 2, 3, 2, 2
arr = np.arange(a1*a2*a3*a4).reshape((a1, a2, a3, a4))

# 拼接操作
result = np.concatenate([arr[:, i] for i in range(a2)], axis=-1)

# 验证结果
print(result.shape)  # 输出 (2, 2, 6)
print(result[0, 0])  # 输出 [0 1 4 5 8 9]，符合预期

方法二：转置+重塑（更高效） #

通过调整轴的顺序，将矩阵数维度移到行维度之后，再直接重塑数组，避免循环和额外数据复制：

import numpy as np

# 示例数组
a1, a2, a3, a4 = 2, 3, 2, 2
arr = np.arange(a1*a2*a3*a4).reshape((a1, a2, a3, a4))

# 转置轴顺序为 (批次, 行, 矩阵数, 列)，再重塑
result = arr.transpose(0, 2, 1, 3).reshape(a1, a3, -1)

# 验证结果
print(result.shape)  # 输出 (2, 2, 6)
print(result[0, 0])  # 输出 [0 1 4 5 8 9]，符合预期

原理说明 #

• 原数组轴顺序：(批次, 矩阵数, 行, 列) → (a1,a2,a3,a4)
• 转置后轴顺序：(批次, 行, 矩阵数, 列) → (a1,a3,a2,a4)，此时每个行对应的多个矩阵的列可以直接合并
• 重塑时用-1让NumPy自动计算合并后的列数（即a2*a4），最终得到目标形状(a1,a3,a2*a4)

对于你提到的原始需求（输入形状(2,3,4,5)），只需替换对应参数，两种方法都能得到(2,4,15)的输出数组。

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