NumPy如何向量化实现索引矩阵到对应列表元素矩阵的转换

核心实现（NumPy环境，完全向量化） #

直接利用NumPy的高级索引特性，无需手动编写循环，输入任意维度的索引矩阵都可以直接输出相同维度的结果矩阵，示例代码如下：

import numpy as np

# 输入参数
z = [[0, 1], [1, 0]] # 任意维度的索引矩阵
lookup_list = [20, 10] # 对应查找列表

# 向量化转换操作
z_np = np.array(z)
lookup_np = np.array(lookup_list)
result_np = lookup_np[z_np]

# 若需要转回Python原生列表格式
result = result_np.tolist()

运行上述代码得到的result即为示例要求的[[20, 10], [10, 20]]。

特性说明 #

• 维度自动适配：无论输入的索引矩阵是2维、3维甚至更高维度，输出结果的维度和索引矩阵完全一致
• 性能优势：操作底层为C实现，相比Python层级的遍历/列表推导，大尺寸矩阵下性能提升可达数百倍
• 其他框架适配：如果使用PyTorch/TensorFlow等深度学习框架，操作逻辑完全一致，直接用查找张量索引索引张量即可。

注意事项 #

需要保证索引矩阵内的所有元素取值范围在[0, len(lookup_list)-1]之间，否则会触发索引越界错误。

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