efficient processing of corresponding elements in two numpy arrays

Question

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.

Are A and B np arrays or lists?

hpaulj
– hpaulj

2016-11-10 23:08:33 +00:00
Commented Nov 10, 2016 at 23:08 — hpaulj
– hpaulj, Commented Nov 10, 2016 at 23:08

Divakar · Accepted Answer · 2016-11-10 21:55:02Z

2

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.

edited Nov 10, 2016 at 21:55

answered Nov 10, 2016 at 21:43

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

efficient processing of corresponding elements in two numpy arrays

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related