如何高效获取N维数组锁定N-1维后的所有子数组组合？

针对任意2-5维数组，高效提取指定维度的子数组集合，可通过维度重排+重塑形状实现，全程基于数组视图操作，无额外数据复制，效率拉满：

核心思路 #

对每个目标维度dim_idx：

1. 使用np.moveaxis将该维度移到数组的最后一位
2. 将数组重塑为(-1, original_shape[dim_idx])，其中-1会自动计算其他维度的总元素数
3. 重塑后的数组每一行就是一个目标子数组，数量恰好等于其他维度的乘积

代码示例（以3维数组为例） #

import numpy as np

# 生成测试数组，shape=(10,20,30)
arr = np.random.rand(10,20,30)

# 遍历所有维度，提取对应子数组集合
for dim_idx in range(arr.ndim):
    # 移动目标维度到最后
    shifted_arr = np.moveaxis(arr, dim_idx, -1)
    # 重塑形状，得到(N, target_dim_size)的数组，N是其他维度的乘积
    sub_arrays = shifted_arr.reshape(-1, arr.shape[dim_idx])
    # 输出结果信息
    print(f"维度{dim_idx}对应的子数组集合：")
    print(f"  原维度大小：{arr.shape[dim_idx]}")
    print(f"  子数组数量：{sub_arrays.shape[0]}")
    print(f"  单个子数组shape：{sub_arrays[0].shape}\n")

适配2-5维数组的通用逻辑 #

上述代码无需修改即可直接适配2-5维数组：

• 对2维数组(a,b)，提取维度0的子数组会得到(b,a)的数组（每个行是shape(a)的子数组），维度1则得到(a,b)的数组（每个行是shape(b)的子数组）
• 对5维数组(d0,d1,d2,d3,d4)，提取维度2时，会得到(d0*d1*d3*d4, d2)的数组，共d0*d1*d3*d4个shape(d2)的子数组

为什么transpose/rollaxis无效？ #

你之前尝试的transpose需要手动指定所有维度的顺序，对动态维度数场景不友好；rollaxis是将维度滚动到首位，后续仍需手动计算重塑，而moveaxis可以直接指定目标位置，配合reshape(-1, ...)完全自动适配任意维度数，逻辑更简洁。

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