如何使用numpy中shape为(x,y)的索引数组提取三维arr对应z轴位置的值

可行实现方案 #

方案1：使用np.take_along_axis（最简洁，适合新手） #

该方法是Numpy专门为这类「沿指定轴、按匹配维度的索引提取元素」的场景设计，不需要理解复杂的索引规则：

import numpy as np
# 示例输入数据
arr = np.arange(99).reshape(3,3,11)
indexes = np.asarray([
    [0,2,2],
    [1,2,3],
    [3,2,10]
])
# 核心实现代码
result = np.take_along_axis(arr, indexes[..., np.newaxis], axis=2).squeeze(axis=2)

输出的result和你给出的expected_result完全一致。
说明：因为take_along_axis要求索引数组和原数组维度数相同，所以用[..., np.newaxis]给indexes加一个长度为1的第三维，提取完成后用squeeze去掉多余的维度即可。

方案2：原生高级索引（性能最优，无需调用额外函数） #

利用Numpy的整数数组索引广播规则，手动构造前两个维度的对应索引即可：

# 生成前两个维度的坐标网格，shape和indexes完全一致
x_idx, y_idx = np.indices(indexes.shape)
# 按位置一一对应提取元素
result = arr[x_idx, y_idx, indexes]

该写法运行效率更高，直接输出shape为(x,y)的目标数组。

原错误写法原因说明 #

• 使用np.take(arr, indexes, axis=2)时，take会把indexes的每个元素都作为第三维索引，对原数组的每个(x,y)位置都返回完整indexes长度的结果，所以输出shape变成了(x,y,x,y)
• arr[indexes, indexes,:]的写法是把indexes同时作为第一、第二维的索引，触发了广播配对逻辑，和「每个位置对应自身x、y坐标」的需求不符

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