如何用两个索引数组实现numpy高维数组的指定元素提取

直接使用NumPy的广播高级索引即可实现，思路是将H、W维度的索引数组reshape为与原数组前3个维度匹配的形状，再对应维度取值：

import numpy as np

# 构造原数组
arr = np.array([[[[[ 7.,  9.],
          [10., 11.]],
         [[19., 18.],
          [20., 16.]]]],
       [[[[24.,  5.],
          [ 6., 10.]],
         [[18., 11.],
          [45., 12.]]]]])

# 定义H、W维度的索引
idx_h = np.array([1, 1, 0, 1]) # 沿H维度的索引
idx_w = np.array([1, 0, 0, 0]) # 沿W维度的索引

# 将索引reshape为与原数组前3个维度(2,1,2)对齐的形状
idx_h = idx_h.reshape(2, 1, 2)
idx_w = idx_w.reshape(2, 1, 2)

# 自动生成前3个维度的广播索引，避免硬编码
idx0, idx1, idx2 = np.indices(arr.shape[:3])

# 高级索引取值
result = arr[idx0, idx1, idx2, idx_h, idx_w]

运行后输出的result就是要求的形状为(2,1,2)的数组：

array([[[11., 20.]],
       [[24., 45.]]])

内容的提问来源于stack exchange，提问作者Sam-gege