多维NumPy数组多元素索引异常问题及正确实现方法

为什么结果不符合预期？ #

你直接使用a[idx]时，NumPy的索引逻辑和你想的不一样：

• idx是形状为(2,3)的二维数组，NumPy会将其视为对a第一个维度的批量索引——把idx里的每个元素都当作a第一维的下标，同时保留a剩余的两个维度。
• 这相当于用idx中的6个元素分别索引a的第一维，每个索引对应一个(3,3)的子数组，最终拼接出形状为(2,3,3,3)的结果，完全不是按三维坐标提取单个元素的逻辑。

如何修改代码得到预期结果？ #

要按三维坐标提取元素，需要让NumPy识别出每个坐标对应的三个维度索引，以下是几种可行方法：

方法1：拆分维度索引 #

将idx的每一列分别作为a三个维度的索引：

import numpy as np

a = np.random.random((3, 3, 3))
idx = np.asarray([[0, 0, 0], [0, 1, 2]])

# 提取每个维度的索引数组
dim1_idx = idx[:, 0]
dim2_idx = idx[:, 1]
dim3_idx = idx[:, 2]

b = a[dim1_idx, dim2_idx, dim3_idx]
print(b.shape)  # 输出 (2,)

方法2：转置坐标数组并转为元组 #

NumPy的多维索引支持元组形式的维度索引，将idx转置后，每一行对应一个维度的所有索引，再转为元组即可：

import numpy as np

a = np.random.random((3, 3, 3))
idx = np.asarray([[0, 0, 0], [0, 1, 2]])

b = a[tuple(idx.T)]
print(b.shape)  # 输出 (2,)

方法3：使用np.take_along_axis #

通过扩展坐标数组的维度，配合take_along_axis实现按坐标提取：

import numpy as np

a = np.random.random((3, 3, 3))
idx = np.asarray([[0, 0, 0], [0, 1, 2]])

# 扩展维度以匹配原数组的轴
idx_expanded = idx[:, np.newaxis, :]
b = np.take_along_axis(a, idx_expanded, axis=0).squeeze()
print(b.shape)  # 输出 (2,)

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