在两个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以矢量化的方式做事。我们首先创建对应于条件的布尔数组。然后，我们使用这些掩码使用boolean-indexing将A索引到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。