2つのnumpy配列の対応する要素の効率的な処理

Question

私は同じサイズの2つの配列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をベクター化して使用することができます。まず、条件に対応するブール値の配列を作成します。その後、我々はpromotesとdemotes取得するboolean-indexingとAへのインデックスにそれらのマスクを使用する -

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となります。