如何按指定索引数组从二维NumPy数组提取生成一维数组

你的需求是从形状为(5, 5527)的二维数组中，按(5527,)的索引数组提取每一列对应行的元素，最终得到一维数组。用np.take出错是因为默认用法不符合你的索引逻辑，试试下面两种简单方法：

方法一：高级索引配对 #

直接用行索引数组和列索引数组进行配对索引，列索引可以用np.arange生成：

import numpy as np

# 假设你的二维数组叫arr，索引数组叫indices
cols = np.arange(arr.shape[1])  # 生成0到5526的列索引
result = arr[indices, cols]

这个方法会逐个取arr[indices[0], 0]、arr[indices[1], 1]……直到最后一列，正好得到你要的(5527,)形状的一维数组。

方法二：用np.choose简化代码 #

np.choose可以直接根据索引从多维数组的指定轴选择元素，这里直接传入索引数组和原数组即可：

result = np.choose(indices, arr)

这个方法更简洁，效果和上面完全一致。

举个小例子验证：

arr = np.array([[1,2,3],[4,1,5],[2,3,1]])  # 形状(3,3)
indices = np.array([0,1,2])
print(arr[indices, np.arange(3)])  # 输出 [1 1 1]
print(np.choose(indices, arr))     # 输出 [1 1 1]

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