如何有效地从 numpy 数组中提取由索引给出的元素列表？答案

【问题标题】：How to extract a list of elements given by their indices from a numpy array efficiently?如何有效地从 numpy 数组中提取由索引给出的元素列表？
【发布时间】：2015-05-12 04:55:10
【问题描述】：

我有一个多维 numpy 数组，我想利用它的一些元素来构造一个一维数组。我需要采取的元素由它们的索引给出，例如：

inds = [(0,0), (0,1), (1,1), (1,0), (0,2)]

我用直截了当的方式解决它：

ls = [] 
for i, j in inds:
   ls += [a[i,j]]

它给出了想要的结果。但是，我已经意识到这个解决方案对于我的目的来说太慢了。是否有可能以更有效的方式做同样的事情？

【问题讨论】：

标签： python numpy optimization

【解决方案1】：

numpy 数组可以用序列索引（更一般地说，numpy 数组）。

例如，这是我的数组a

In [19]: a
Out[19]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

i 和j 保存inds 数组的第一个和第二个坐标的序列：

In [20]: i
Out[20]: [0, 0, 1, 1, 0]

In [21]: j
Out[21]: [0, 1, 1, 0, 2]

您可以使用这些从a 中提取相应的值：

In [22]: a[i, j]
Out[22]: array([0, 1, 6, 5, 2])

如果您的代码中已经有inds，您可以使用zip 将元组列表分成i 和j：

In [23]: inds
Out[23]: [(0, 0), (0, 1), (1, 1), (1, 0), (0, 2)]

In [24]: i, j = zip(*inds)

In [25]: i
Out[25]: (0, 0, 1, 1, 0)

In [26]: j
Out[26]: (0, 1, 1, 0, 2)

或者，如果inds 是一个形状为 (n, 2) 的数组，如下所示：

In [27]: inds = np.array(inds)

In [28]: inds
Out[28]: 
array([[0, 0],
       [0, 1],
       [1, 1],
       [1, 0],
       [0, 2]])

您可以简单地将inds 的转置分配给i, j：

In [33]: i, j = inds.T

In [34]: i
Out[34]: array([0, 0, 1, 1, 0])

In [35]: j
Out[35]: array([0, 1, 1, 0, 2])

In [36]: a[i, j]
Out[36]: array([0, 1, 6, 5, 2])

【讨论】：