如何高效构建由多个(k,m)矩阵组成的(n*k,m)大小的NumPy数组

需求说明 #

现有三个形状为(2,4)的NumPy数组A、B、C，需要构建一个形状为(6,4)的数组X，要求X的行按A的第1行 → B的第1行 → C的第1行 → A的第2行 → B的第2行 → C的第2行的顺序循环排列。原循环实现效率低且无法适配数组数量变化的场景，需要通用且高效的矢量化实现方案。

原数组定义 #

import numpy as np

A = np.array([[0, 2, 1, 2],
              [3, 1, 4, 3]])
B = np.array([[1, 2, 1, 0],
              [0, 4, 3, 1]])
C = np.array([[0, 4, 3, 2],
              [3, 0, 1, 0]])

高效通用实现方案 #

利用NumPy的矢量化操作，通过stack+swapaxes+reshape组合实现，完全避免循环，且适配任意数量的同形状数组：

# 将所有目标数组堆叠为三维数组，调整轴顺序后展平
X = np.stack([A, B, C]).swapaxes(0, 1).reshape(-1, 4)

代码解释

1. np.stack([A, B, C])：将三个数组在第0轴堆叠，得到形状为(3, 2, 4)的三维数组（第0轴对应数组A/B/C，第1轴对应原数组的行，第2轴对应列）；
2. .swapaxes(0, 1)：交换第0轴和第1轴，得到形状为(2, 3, 4)的数组（第0轴对应原数组的行，第1轴对应数组A/B/C）；
3. .reshape(-1, 4)：将前两个轴合并为一个轴，自动计算长度为2*3=6，最终得到形状(6,4)的目标数组X。

验证结果

打印X可得到符合要求的结果：

print(X)
# 输出：
[[0 2 1 2]
 [1 2 1 0]
 [0 4 3 2]
 [3 1 4 3]
 [0 4 3 1]
 [3 0 1 0]]

优势说明 #

• 高效性：完全基于NumPy底层矢量化操作，比Python循环快几个数量级，尤其适合大数组场景；
• 通用性：只需修改stack的输入列表即可适配任意数量的同形状数组（比如新增数组D/E，只需改为np.stack([A,B,C,D,E])）。

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