如何用Numpy向量化实现通过掩码提取二维数组指定元素？

现有两个Numpy数组：

x     # 形状为(n, m)
mask  # 形状为(n)，每个元素取值范围是0到m-1

目标是提取x中对应位置的元素，得到形状为n的数组，满足：

out[i] = x[i, mask[i]]

用for循环的实现方式如下：

import numpy as np
out = np.zeros(n)
for i in range(n):
    out[i] = x[i, mask[i]]

以下是几种高效的Numpy向量化实现方案：

• 方案一：直接使用配对索引
  这是最简洁直观的方式，利用np.arange(n)生成行索引，和mask的列索引组合，直接定位目标元素：

out = x[np.arange(n), mask]

Numpy会自动将两个一维索引数组广播为匹配的索引对，精准提取每个(i, mask[i])位置的元素。

• 方案二：使用take_along_axis函数
  通过take_along_axis可以沿着指定轴提取对应索引的元素，需要先将mask调整为和x列维度匹配的形状：

out = np.take_along_axis(x, mask[:, np.newaxis], axis=1).squeeze()

mask[:, np.newaxis]将mask从(n,)转换为(n,1)，与x的列维度对齐，提取后用squeeze()去除多余维度，得到(n,)形状的结果。

• 方案三：扁平化数组后计算索引
  先将x扁平化，再通过计算每个目标元素在扁平化数组中的位置来提取：

flat_x = x.flatten()
indices = np.arange(n) * m + mask
out = flat_x[indices]

这种方式适合理解Numpy数组索引的底层逻辑，通过行号乘以列数加上列索引得到扁平化后的位置。

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