如何利用height数组高效将3D NumPy数组切片为2D数组

我有两个NumPy数组：

a.shape
# (1024, 768, 100)
heights.shape
# (1024, 768)

我需要从数组a中提取一个二维结果，沿轴2（第三维）取单个元素，每个位置的索引由heights数组对应坐标的值决定，预期结果形状为(1024, 768)，逻辑上类似：

result = a[:,:,h]  # h为heights对应位置的取值

嵌套循环可以实现但效率极低：

result = np.zeros(heights.shape)
for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        h = heights[i][j]
        result[i][j] = a[i][j][h]

示例 #

import numpy as np

a = np.arange(30).reshape((3,2,5)) + 5
# array([[[ 5,  6,  7,  8,  9],
#         [10, 11, 12, 13, 14]],
#
#        [[15, 16, 17, 18, 19],
#         [20, 21, 22, 23, 24]],
#
#        [[25, 26, 27, 28, 29],
#         [30, 31, 32, 33, 34]]])

heights = np.asarray([[0, 2], [1, 4], [1, 3]])
# array([[0, 2],
#        [1, 4],
#        [1, 3]])

expected_result = np.array([[ 5, 12],
                            [16, 24],
                            [26, 33]])

高效实现方法 #

利用NumPy的高级索引可以直接完成操作，完全避免循环，性能和原生NumPy操作一致：

方法一：使用索引网格

# 生成前两维的索引网格
i, j = np.indices(heights.shape)
# 提取对应元素
result = a[i, j, heights]

方法二：利用广播简化写法

result = a[np.arange(a.shape[0])[:, None], np.arange(a.shape[1]), heights]

验证示例输出：

print(result)
# 输出:
# [[ 5 12]
#  [16 24]
#  [26 33]]

原理说明：
两种方法本质都是构造三维的索引组合，让每个位置(x,y)对应a[x][y][heights[x][y]]。NumPy会自动将这些索引对应到数组元素，直接生成目标二维数组，比循环效率高几个数量级，尤其适合大尺寸数组。

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