如何用索引数组从任意维度NumPy ndarray提取二维子数组？

核心思路 #

要实现任意维度下的动态索引提取，关键是把索引矩阵X转换成NumPy高级索引要求的维度匹配的索引元组：将X转置后，每一行对应dummy前N-1个维度的索引数组，最后再加上一个切片来保留dummy的最后一个维度。

代码实现 #

直接用以下通用写法即可适配所有符合条件的维度：

import numpy as np

# 示例数据
dummy = np.random.rand(5,2,100)
X = np.array([[0,1],[4,1],[2,0]])

# 通用提取逻辑
result = dummy[tuple(X.T) + (slice(None),)]

# 验证结果形状：预期是(3,100)
print(result.shape)  # 输出 (3, 100)

适配更高维度的情况 #

比如当dummy是5×2×4×6×100的多维数组，X是3×4的索引矩阵时，同样适用：

dummy_high_dim = np.random.rand(5,2,4,6,100)
X_high_dim = np.array([[0,1,3,5], [4,0,2,1], [2,1,0,3]])  # 3行4列，对应前4个维度的索引

result_high_dim = dummy_high_dim[tuple(X_high_dim.T) + (slice(None),)]
print(result_high_dim.shape)  # 输出 (3, 100)

为什么之前的方法失效？ #

• dummy[X,:]：把X整体当作第一个维度的索引，相当于对第一个维度做批量索引，而非分别匹配前N-1个维度，导致维度混乱。
• 迭代索引：每次迭代会改变数组的维度结构，后续索引无法对应原数组的维度位置，自然得不到正确结果。

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