无需逐个切片转置，如何批量处理NumPy数组的方形子块？

假设我们有一个由n个k×k方形子块横向拼接而成的NumPy数组A，结构为：

A = [A₁ | A₂ | ... | Aₙ]

其中A的形状为(k, n*k)，现在需要将每个子块转置后再横向拼接，得到：

B = [A₁^T | A₂^T | ... | Aₙ^T]

完全不需要逐个切片子块再转置，可以通过reshape与轴交换的组合操作实现，全程利用NumPy的向量化操作，效率远高于循环处理。

具体实现步骤（附示例） #

import numpy as np

# 定义参数与示例数组
k = 2  # 子块边长
n = 2  # 子块数量
A = np.array([[1, 2, 3, 4],
              [5, 6, 7, 8]])

# 批量转置子块并拼接
B = A.T.reshape(n, k, k).swapaxes(0, 1).reshape(k, n*k)

操作原理拆解 #

1. A.T：将原数组转置，把每个子块的列转换为行，此时A.T的形状为(n*k, k)，前k行对应A₁^T的所有行。
2. reshape(n, k, k)：将转置后的数组重新塑形为(n, k, k)，此时每个索引i对应的子数组arr[i]就是A_{i+1}^T。
3. swapaxes(0, 1)：交换前两个维度，将数组形状变为(k, n, k)，实现把所有子块转置后的行按顺序排列。
4. reshape(k, n*k)：最后重新塑形为目标形状(k, n*k)，得到横向拼接所有转置子块的结果。

验证结果 #

运行上述代码后，B的输出为：

array([[1, 5, 3, 7],
       [2, 6, 4, 8]])

完全符合[A₁^T | A₂^T]的预期结构。

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