有效處理兩個numpy數組中的對應元素

Question

我有兩個相同大小的數組A和B.我需要根據B的項目獲取數組中項目的值和索引。這裏是我的代碼：

promotes = [a for a,b in zip(A, B) if b == 1] 
demotes = [a for a,b in zip(A, B) if b == 0] 
promotes_index = [k for k,(a,b) in enumerate(zip(A, B)) if b == 1] 
demotes_index = [k for k,(a,b) in enumerate(zip(A, B)) if b == 0] 

我相信有一個更有效的方式來計算宣傳，降級和指數。

更簡單地說，如果A和促進就像：

A = array([ 4, 6, 9, 10])) 
promotes = array([4, 9]) 

我怎麼會從A promote_index和促進：

promotes_index = array([0, 2]) 

我明白任何迴應。

Answer 1

我們可以在那裏使用array programming以矢量化的方式進行操作。我們首先創建對應於條件的布爾數組。然後，我們使用這些面具指數爲A與boolean-indexing得到promotes和demotes -

mask1, mask2 = B==1, B==0 
promotes, demotes = A[mask1], A[mask2] 

最後，我們使用np.flatnonzero這些面具來得到相應的指數 -

promotes_index, demotes_index = np.flatnonzero(mask1), np.flatnonzero(mask2) 

請注意，如果B僅由1s和0s組成，mask2將僅僅是~mask1。