python - 使用numpy中的一维数组从二维数组中选择多个元素

Question

我在 numpy 中有两个数组。第一个是二维数组，可以认为是向量列表。第二个是一维数组，可以将其视为二维数组的索引列表。

我想使用一维数组的索引选择二维数组的元素。现在我一直在做

        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 中有没有更好的方法来做到这一点？

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]]