How do I select elements of an two array given condition?

Viewed 144

Suppose I have a 3 numpy arrays

x1 = [5, 2, 3, 1, 4, 5] x2 = [8,14,22,33,0,7] y = [0,0,1,0,1,1]

I want to select the elements in x1 when y==0 and in x2 when y==1 output = [5,2,22,1,0,7]

What is the most efficient way ?

3 Answers

Use numpy.where:

np.where(y, x2, x1)

Output:

array([ 5,  2, 22,  1,  0,  7])

Try masking and simple element-wise addition

import numpy as np

x1 = np.array([5, 2, 3, 1, 4, 5])

x2 = np.array([8,14,22,33,0,7])

y = np.array([0,0,1,0,1,1]).astype(bool)

output = x1*~y + x2*y
array([ 5,  2, 22,  1,  0,  7])

In case of list, instead of np.array, you can use:

[x2[i] if v else x1[i] for i,v in enumerate(y)]
#[5, 2, 22, 1, 0, 7]
Related