将3d数组拆分为具有不同第二索引的2d数组的最pythonic方法是什么？

Question

我有一个函数foo，它返回一个形状为（1000,3,1000）的数组 如何将其拆分为三个数组a（1000,1000）和b（1000,1000）和c（1000,1000） 我正在寻找这样的东西：

a,b,c = foo()[:,*,:]

将导致a，b，c等于：

res = foo()
a = res[:,0,:]
b = res[:,1,:]
c = res[:,2,:]

Answer 1

我相信以下任何一项都有效：

numpy.split(foo(), 3, axis=1)
numpy.hsplit(foo(),3)
list(numpy.swapaxes(foo(), 0, 1))

Answer 2

这也是一件好事：

a,b,c = np.rollaxis(foo(), 1)

Answer 3

与您的previous question类似：

def foo():
    return np.ones((1000, 3, 1000))

a, b, c = [np.squeeze(arr) for arr in np.hsplit(foo(), foo().shape[1])]

for arr in [a, b, c]:
    print(arr.shape)
# (1000, 1000)
# (1000, 1000)
# (1000, 1000)

如果没有np.squeeze，你将获得（1000,1,1000）的形状。