如何从二维NumPy数组按列表选取样本？求更高效向量化实现

没问题！咱们来聊聊怎么用NumPy的向量化风格高效搞定这个需求～

先对应你的预期结果(5, 6, 8)，我先假设一个典型场景例子：比如你的二维矩阵和目标索引列表是这样的：

import numpy as np
matrix = np.array([[1,2,3],
                   [4,5,6],
                   [7,8,9]])
# 要提取的位置索引
indices = [(1,1), (1,2), (2,1)]

我猜你之前可能用了Python循环逐个提取？比如这样：

result = []
for i,j in indices:
    result.append(matrix[i,j])

虽然能得到结果，但面对大矩阵或大量索引时，Python循环的效率就跟不上了——这时候就得用NumPy的向量化操作来提速。

最直观的向量化实现 #

把索引列表拆成行索引数组和列索引数组，直接用这两个数组去索引矩阵就行：

# 拆分行、列索引
rows, cols = zip(*indices)
# 一次性提取所有目标元素
result = matrix[rows, cols]
print(result)  # 输出: [5 6 8]

这种方式完全跳过了Python层面的循环，所有操作都在NumPy底层的C代码里执行，数据量越大，效率优势越明显。

如果索引是NumPy数组的情况 #

要是你的索引本来就是NumPy数组格式（而非Python列表），可以直接切片拆分：

indices = np.array([[1,1], [1,2], [2,1]])
result = matrix[indices[:, 0], indices[:, 1]]

同样是纯向量化操作，效率拉满。

另一种思路：转成扁平索引 #

你也可以把二维索引转换成矩阵的一维扁平索引，再用np.take提取：

# 将二维索引转换为矩阵的扁平索引
flat_indices = np.ravel_multi_index(zip(*indices), matrix.shape)
# 通过扁平索引提取元素
result = np.take(matrix, flat_indices)

这种方法在需要重复使用扁平索引的场景下会很方便。

总的来说，这几种方法都是纯NumPy向量化的实现，比循环高效太多，其中第一种拆分行列的方式最直观也最常用～

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