如何用Numpy数组按行列索引提取矩阵元素并实现高效操作？

要实现从Numpy矩阵中按「第i行取第b[i]列」的规则提取元素，最高效的方式是使用整数数组索引（高级索引）——这是Numpy原生支持的矢量化操作，完全避开Python循环，效率最优。

实现步骤 #

1. 构造行索引数组：对应矩阵的每一行序号，即0到a.shape[0]-1。
2. 直接用a[行索引数组, 列索引数组b]提取目标元素。

代码示例 #

import numpy as np

# 初始化示例矩阵和列索引数组
a = np.tile(np.arange(1920), (41, 1))
b = np.arange(40, -1, -1)

# 生成行索引数组
row_indices = np.arange(a.shape[0])

# 提取目标元素
c = a[row_indices, b]

# 验证结果
print(c)  # 输出：[40 39 38 ...  0]

效率说明 #

这种方式是Numpy底层优化的矢量化操作，所有计算在C层面完成，比手动写Python循环快几个数量级，尤其当矩阵规模较大时优势更明显。

注：你的示例中a的每行都是0~1919，所以c恰好等于b，但上述方法是通用解法，适用于任意结构的矩阵。

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