在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數組選擇多個元素

回答

相關問題