如何从二维NumPy数组中按指定索引高效提取元素？

问题场景 #

给定二维NumPy数组：

import numpy as np

A = np.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])  # Shape: (3, 4)

以及一维索引数组：

B = np.array([0, 2, 1])  # Shape: (3,)

需要提取A的第i行中，对应B[i]指定索引的元素，最终得到结果C = [0, 6, 9]，要求替代for循环实现更高效的向量化操作。

优化实现方法 #

方法1：直接利用二维数组索引机制

这是最简洁高效的方式，通过生成行索引数组，与B组合成索引对提取元素：

row_indices = np.arange(A.shape[0])
C = A[row_indices, B]
print(C)  # 输出：array([0, 6, 9])

原理：row_indices生成与A行数匹配的序列[0,1,2]，与B的[0,2,1]一一对应，形成索引对(0,0)、(1,2)、(2,1)，直接定位到目标元素。

方法2：使用np.take_along_axis

如果需要更明确的轴方向控制，可以用np.take_along_axis，需先将B转换为二维数组匹配轴维度：

# 将B转换为(3,1)的二维数组，匹配A的列轴维度
B_2d = B[:, np.newaxis]
C = np.take_along_axis(A, B_2d, axis=1).flatten()
print(C)  # 输出：array([0, 6, 9])

原理：take_along_axis会沿着指定轴（这里是列轴axis=1），在每一行取对应B_2d位置的元素，最后通过flatten()转换为一维数组。

注意事项 #

• 确保A和B都是NumPy数组（而非Python列表），否则无法使用上述向量化操作。
• 这些方法都是NumPy内置的向量化实现，相比for循环，在数组规模较大时性能提升显著。

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