使用NumPy中的一个一维数组选取一个二维数组中的多个元素

Question

我在numpy中有两个数组。第一个是2D数组，可以看作是向量列表。第二个是1D数组，可以看作是对2D数组的索引列表。

我想使用1D数组的索引选择2D数组的元素。目前我一直在做以下操作：

        z=rnd.rand(2,10) # a list of 2d vectors of length 10
        z_idx=rnd.randint(2,size=z.shape[1]) #indices selecting a dimension of the 2d vector

        result=np.array([z[z_idx[i],i] for i in xrange(len(z_idx))])

但这样做非常缓慢。

在numpy中有更好的方法吗？

- user1149913

---

在numpy中，您可以使用(bool)掩码选择数组的元素。您查看过吗？一般来说，您可以从1D数组创建此类遮罩，然后将此遮罩应用于2D数组。 - Dr. Jan-Philip Gehrcke

Answer 1

可能是最简单的方法：

result = z[z_idx].diagonal()

也许更有效率的方法是使用 arange:

result = z[z_idx, np.arange(z_idx.size)]

更恰当而等效的是np.indices ：

result = z[z_idx, np.indices(z_idx.shape)[0]]

- ecatmur

---

请注意，第一个方法将创建一个(len(z_idx), len(z_idx))的临时数组，这可能会成为一个问题。 - Pierre GM