如何用形状为(h,w)的数组Y索引形状为(h,w,d)的数组X的最后维度？

针对你提出的需求——从形状为(h, w, d)的数组X中，用形状为(h, w)的索引数组Y提取最后一维对应位置的元素，最终得到(h, w)的结果，下面是几种实用的numpy实现方法：

1. 花式索引（最直观的写法） #

通过生成和Y匹配的前两维索引，组合Y作为最后一维索引，直接定位目标元素：

import numpy as np

# 构造示例数据
h, w, d = 2, 3, 4
X = np.random.rand(h, w, d)
Y = np.random.randint(0, d, size=(h, w))

# 生成行、列索引矩阵
row_idx, col_idx = np.indices((h, w))
# 提取对应元素
output = X[row_idx, col_idx, Y]

这里row_idx和col_idx分别是(h, w)形状的矩阵，对应每个位置的行、列坐标，加上Y的索引，就能精准取到X中每个(h,w)位置对应的最后一维元素，结果自然是(h,w)形状。

2. 使用np.take函数（简洁高效） #

np.take可以直接沿指定维度提取元素，指定axis=2（最后一维）后传入Y，再去掉冗余的单维度：

output = np.take(X, Y, axis=2).squeeze()

np.take返回的结果形状是(h, w, 1)，用squeeze()去掉最后一个长度为1的维度，就得到了目标的(h,w)数组。

3. 重塑数组+扁平化索引（适合理解底层逻辑） #

先把X和Y都扁平化，提取后再恢复原形状：

# 把X重塑为二维数组，Y重塑为一维数组
X_flat = X.reshape(-1, d)
Y_flat = Y.reshape(-1)
# 提取元素后再重塑回原形状
output = X_flat[np.arange(h*w), Y_flat].reshape(h, w)

np.arange(h*w)生成扁平化后的行索引，搭配Y_flat的列索引，就能准确取到每个位置的元素，最后重塑回(h,w)即可。

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