如何用Numpy操作根据二维索引数组提取二维ndarray的指定列元素？

你遇到的问题是因为直接使用arr[arr_idx]时，Numpy会将arr_idx中的每个元素视为行索引去提取整行，因此得到三维数组。要实现按行提取对应列的需求，需要构造与arr_idx匹配的行索引数组，结合列索引一起定位元素。

实现代码 #

import numpy as np

arr = np.array([[1,2,3],[4,5,6],[7,8,9]])
arr_idx = np.array([[0,2],[1,2],[2,1]])

# 构造与arr_idx形状匹配的行索引数组
row_indices = np.arange(arr.shape[0])[:, np.newaxis]
# 同时使用行、列索引提取元素
result = arr[row_indices, arr_idx]

print(result)
# 输出结果：
# [[1 3]
#  [5 6]
#  [9 8]]

原理说明 #

• np.arange(arr.shape[0])生成行号数组[0,1,2]，通过[:, np.newaxis]将其转换为形状为(3,1)的列向量。
• 这个列向量会和形状为(3,2)的arr_idx自动广播，形成形状为(3,2)的行索引矩阵，每个位置对应arr的行号，arr_idx对应列号。
• 这种索引方式是完全向量化的，避免了循环，处理大数据量时效率远高于列表推导式。

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