拆分NumPy数组时如何直接实现维度缩减？

问题场景 #

给定代码：

import numpy as np

a = np.arange(8).reshape(2, 2, 2)
b = np.split(a, 2)
print(b[0].shape)

运行结果为：

(1, 2, 2)

需要将数组a拆分为形状为(2, 2)的子数组列表，且无需额外步骤去除多余维度。

解决方案 #

直接迭代原数组即可——NumPy数组在迭代时会自动沿第一个维度拆分，同时自动去掉该维度：

import numpy as np

a = np.arange(8).reshape(2, 2, 2)
b = list(a)
print(b[0].shape)  # 输出 (2, 2)

如果觉得转列表麻烦，手动索引提取也能达到效果：

b = [a[0], a[1]]

要是一定要用np.split相关方法，也可以一步挤压掉多余维度，不过不如上面的方法简洁：

b = [sub.squeeze() for sub in np.split(a, 2)]

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