在numpy中使用1d数组选择多个元素

我有numpy中的两个数组。第一个是2d数组，可以认为它是一个向量列表。第二个是1d数组，可以认为它是2d数组中的索引列表。在numpy中使用1d数组选择多个元素

我想使用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中做到这一点？

来源

2012-08-31 user1149913

在numpy的，您可以使用（布尔）蒙版选择数组的元素。你看过吗？一般来说，您可以从您的一维阵列中创建这样的遮罩，然后将此遮罩应用于2D阵列。 –

也许最简单的方法：

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

来源

2012-08-31 13:28:41 ecatmur

请注意，第一种方法将创建一个'（len（z_idx），len（z_idx））'临时数组，这可能是一个问题。 –

很好！谢谢 – user1149913

在numpy中使用1d数组选择多个元素

回答

相关问题