如何使用NumPy根据行、列索引列表快速提取二维数组指定元素

现有实现与痛点 #

需要从二维NumPy数组中，提取给定行索引列表、列索引列表笛卡尔积对应的所有元素，当前采用Python嵌套循环实现：

import numpy as np
z = np.arange(15).reshape(3,5)
indexx = [0,2]
indexy = [1,2,3,4]

zz = []
for i in indexx:
  for j in indexy:
    zz.append(z[i][j])

运行输出结果为：

zz >> [1, 2, 3, 4, 11, 12, 13, 14]

这种实现虽然结果正确，但嵌套循环是Python解释层的逐元素迭代，当数组规模大、索引列表较长时，解释执行的额外开销会导致运行速度非常慢。

高效向量化实现 #

完全可以借助NumPy原生的广播和向量化索引能力，跳过Python层循环，用底层C实现的逻辑完成计算，性能可以提升数个数量级。

实现思路 #

1. 调整行索引的形状为列向量（增加一维），利用NumPy的广播规则，自动生成覆盖所有行、列索引笛卡尔积的二维索引矩阵
2. 直接用这组二维索引矩阵对原数组做花式索引，得到匹配位置的元素矩阵，再展平为一维列表即可得到和原有逻辑完全一致的输出

示例代码 #

# 行索引转列向量，形状从 (2,) 变为 (2, 1)：[[0], [2]]
row_idx = np.array(indexx)[:, None]
# 列索引保持一维行向量形状：[1, 2, 3, 4]
col_idx = np.array(indexy)

# 广播自动匹配为形状 (2,4) 的索引对，直接提取元素后展平
zz = z[row_idx, col_idx].ravel().tolist()

运行上述代码得到的输出和原嵌套循环结果完全一致，没有任何Python层的循环逻辑，哪怕索引列表长度达到十万、百万级，也能快速完成计算。如果后续不需要把结果转成Python列表，直接保留NumPy数组类型，还可以省去类型转换开销，性能进一步提升。

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