如何使用Numpy实现索引数组对三维lookup数组的多维取值操作？

你可以通过Numpy原生的高级索引机制实现需求，完全不需要Python层面的循环，性能远高于手写循环实现：

实现代码 #

import numpy as np

# 你的原始数组定义
lookup = np.array([
    [["a", "b", "c"], ["d", "e", "f"]],
    [["g", "h", "i"], ["j", "k", "l"]],
    [["m", "n", "o"], ["p", "q", "r"]],
])
i = np.array([
    [[0,0,0], [1,0,2],],
    [[2,0,1], [1,0,2],],
])

# 核心索引操作：拆分i的最后一维对应lookup的三个维度索引
out = lookup[i[..., 0], i[..., 1], i[..., 2]]

执行后打印out即可得到你预期的结果：

array([['a', 'i'],
       ['n', 'i']], dtype='<U1')

通用兼容写法 #

如果需要适配不同维度的lookup数组（只要i的最后一维长度等于lookup的维度数），可以用更通用的np.take_along_axis实现：

out = np.take_along_axis(lookup[tuple([np.newaxis] * i.ndim)], i, axis=-1).squeeze(-1)

原理解释 #

• 第一种写法直接拆分索引数组i的最后一维，分别作为lookup三个维度的索引，Numpy的广播机制会自动匹配前两个维度的形状，直接返回和i前两维形状一致的结果。
• 第二种写法先给lookup补全和i相同数量的前导维度完成广播对齐，再沿最后一维取索引值，最后压缩掉多余的维度，无需手动拆分索引，适配性更强。

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