NumPy ndarray索引后形状不符，如何获取(10,3)的目标形状？

问题场景 #

运行以下代码后，得到的数组形状是(3,10)，但预期是(10,3)：

x = np.ones((10,10,10))
x = x[2,:,[2,3,7]]
print(x.shape)

解决方法 #

方法1：分步索引 #

先提取前两个维度的目标部分，再索引第三维度（此时已变为二维数组的第二维度）：

x = np.ones((10,10,10))
x = x[2,:,:][:,[2,3,7]]
print(x.shape)  # 输出 (10,3)

方法2：转置调整维度顺序 #

直接对索引后的数组进行转置操作，快速调换维度顺序：

x = np.ones((10,10,10))
x = x[2,:,[2,3,7]].T
print(x.shape)  # 输出 (10,3)

方法3：使用np.take指定轴 #

通过np.take函数明确指定索引的轴，避免维度顺序混乱：

import numpy as np
x = np.ones((10,10,10))
x = np.take(x[2,:,:], [2,3,7], axis=1)
print(x.shape)  # 输出 (10,3)

原因说明 #

当在三维数组中混合使用单元素索引（第一维度的2）、切片（第二维度的:）和列表索引（第三维度的[2,3,7]）时，NumPy的索引规则会将列表索引对应的维度前置，导致最终形状变为(3,10)。通过上述方法可以调整维度顺序，得到预期的(10,3)。

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