如何利用二维索引数组从三维数组中提取目标二维数组？

可以利用NumPy的高级索引实现需求，以下是两种常用的简洁方法：

• 方式一：使用np.take_along_axis（推荐，可读性强）
  该函数专门用于沿指定轴按索引提取元素，步骤如下：

  import numpy as np
  
  # 创建示例数组
  data_3d = np.random.random((10, 33, 66))
  index_2d = np.random.randint(0, 10, size=(33, 66))
  
  # 扩展索引数组的维度，匹配原三维数组的轴结构
  index_expanded = index_2d[np.newaxis, ...]
  # 沿第一维度提取元素后，去除多余的单维度
  result = np.take_along_axis(data_3d, index_expanded, axis=0).squeeze()
  

• 方式二：构造广播索引
  通过生成对应维度的网格索引，结合广播机制定位元素：

  import numpy as np
  
  data_3d = np.random.random((10, 33, 66))
  index_2d = np.random.randint(0, 10, size=(33, 66))
  
  # 生成第二、第三维度的网格索引
  idx_2, idx_3 = np.meshgrid(np.arange(33), np.arange(66), indexing='ij')
  # 组合索引提取目标元素
  result = data_3d[index_2d, idx_2, idx_3]
  

验证结果形状：

print(result.shape)  # 输出 (33, 66)

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