python - 重新排序 numpy ndarray 的最后一个维度

Question

我有一个未指定维度的数组(可能是 1D、2D、3D、4D，...)。我想使用与最后一个维度大小相同的索引数组来申请对最后一个维度进行重新排序。我知道如何使用虚拟 if ... if 语句来做到这一点:

import numpy as np

a = np.ones((10, 10, 20, 40,30,10))
index=np.arange(a.shape[-1])

if len(a.shape)==1:
    a=a[index]
elif len(a.shape)==2:
    a=a[:,index]
elif len(a.shape)==3:
    a=a[:,:,index]
elif len(a.shape)==4:
    a=a[:,:,:,index]
elif  len(a.shape)==5:
    a=a[:,:,:,:,index]
else:
    print('Dimensions too high!!!')

我当然对此不满意，因为它不能推广到任何维度。如果有人能向我指出正确的方法，我将不胜感激!

谢谢!

Answer 1

只需使用省略号运算符即可实现通用n-dim索引 -

a[...,index]

More info on ellipsis when indexing into arrays 。

Extensive Q&A on ellipsis .