CODEWITHSUNDEEP
pythonarraysnumpy
crownLearning
crownLearning

Reputation: 35

efficient processing of corresponding elements in two numpy arrays

I have two arrays A and B of the same size. I'd need to get the values and indices of items in the arrays according to B's items. Here is my code:

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]

I'm sure there is a more efficient way to compute promotes, demotes, and the indices.

In more simple words, if A and promotes are like:

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

how could I get promote_index from A and promotes:

promotes_index = array([0, 2])

I appreciate any response.

Upvotes: 0

Views: 56

Answers (1)

Divakar
Divakar

Reputation: 221574

We could use array programming there to do things in a vectorized manner. We start off by creating boolean arrays corresponding to the conditionals. Then, we use those masks to index into A with boolean-indexing to get promotes and demotes -

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

Finally, we use np.flatnonzero on those masks to get the corresponding indices -

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

Please note that if B consists of 1s and 0s only, mask2 would be simply ~mask1.

Upvotes: 2

Related Questions