numpy切片奇怪的行为



我有一个像这样的5维数组

a=np.random.randint(10,size=[2,3,4,5,600])
a.shape #(2,3,4,5,600)

我想获得第二维的第一个元素,以及最后一个维度的几个元素

b=a[:,0,:,:,[1,3,5,30,17,24,30,100,120]]
b.shape #(9,2,4,5)
如您所见,最后一个维度被自动转换为第一个维度。为什么?如何避免这种情况呢?

这个行为在numpy文档中有描述。在表达式

a[:,0,:,:,[1,3,5,30,17,24,30,100,120]]

0[1,3,5,30,17,24,30,100,120]都是高级索引,用切片分隔。正如文档所解释的,在这种情况下,来自高级索引的维度将首先出现在结果数组中。

如果我们用切片0:1替换0,它将改变这种情况(因为它将只留下一个高级索引),然后维度的顺序将被保留。因此,解决此问题的一种方法是使用0:1切片,然后挤压适当的轴:

a[:,0:1,:,:,[1,3,5,30,17,24,30,100,120]].squeeze(axis=1)

或者,可以保留两个高级索引,然后重新排列轴:

np.moveaxis(a[:,0,:,:,[1,3,5,30,17,24,30,100,120]], 0, -1)

最新更新