如何使用Numpy根据坐标列表从三维数组中提取对应元素且避免循环

你需要用Numpy的高级整数索引实现需求，有两种常用的写法：

方法1：直接拆分坐标维度（适合固定维度场景） #

把second数组的每一列分别作为first三个维度的索引即可，代码如下：

import numpy as np

# 示例数据
first = np.array([[[100, 101, 102],
                  [103, 104, 105],
                  [106, 107, 108]],

                 [[109, 110, 111],
                  [112, 113, 114],
                  [115, 116, 117]],

                 [[118, 119, 120],
                  [121, 122, 123],
                  [124, 125, 126]]])

second = np.array([[0, 1, 0],
                   [1, 1, 0],
                   [1, 0, 0],
                   [0, 0, 0],
                   [0, 1, 1],
                   [1, 1, 1],
                   [1, 0, 1],
                   [0, 0, 1]])

# 核心实现代码
result = first[second[:, 0], second[:, 1], second[:, 2]]

运行后得到的result和预期结果完全一致：array([103, 112, 109, 100, 104, 113, 110, 101])。

方法2：通用写法（适配任意维度数组） #

如果数组维度不固定，可以用转置后转元组的写法，无需手动拆分每一列：

result = first[tuple(second.T)]

原理解释 #

直接写first[second]结果不对的原因是：Numpy会把传入的二维数组second默认作为第一个维度的索引序列，返回维度为(8, 3, 3, 3)的数组，和预期不符。而我们把second的每一列单独作为对应维度的索引，Numpy会按行匹配三个维度的坐标值，逐行取出对应位置的元素，得到一维的结果数组，全程都是向量化操作，没有循环开销。

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