如何以最简方式实现三维数组按二维索引数组取值？及对应索引指南章节

最简实现方式 #

在Python的NumPy库中，可通过整数数组索引快速实现需求：

import numpy as np

# 示例初始化（替换为你的实际数组即可）
M, N, K = 3, 4, 5
a = np.random.rand(M, N, K)
b = np.random.randint(0, K, size=(M, N))

# 生成行列索引矩阵，配合b完成元素提取
rows, cols = np.indices(b.shape)
c = a[rows, cols, b]

如果追求极致简洁，也可以写成一行：

c = a[tuple(np.indices(b.shape)) + (b,)]

验证：最终c的每个元素c[i,j]都严格等于a[i,j,b[i,j]]，完全符合要求。

对应索引指南的部分 #

该操作属于NumPy索引体系中的整数数组索引（归属于高级索引范畴）。这类索引允许用数组作为索引参数，按元素级别的对应关系从原数组中选取元素，恰好适配此处"每个(i,j)位置对应第三个维度的b[i,j]索引"的场景。

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