CODEWITHSUNDEEP
pythonarraysnumpyscipyextract
Vincent
Vincent

Reputation: 60361

How to select the rows where a condition is satisfied

Consider three numpy arrays:

x = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) 
y = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
z = np.array([0, 0, 1, 1, 0, 0, 1, 1, 0, 0])

How can I produce an array from x, where a condition is satisfied involving y and z. For example "extract x where y+z = 2" would return:

np.array([3, 7])

(as I have a lot of data, I want this to be as fast as possible)

Upvotes: 3

Views: 76

Answers (1)

falsetru
falsetru

Reputation: 368954

>>> x = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> y = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
>>> z = np.array([0, 0, 1, 1, 0, 0, 1, 1, 0, 0])

>>> y + z == 2
array([False, False, False,  True, False, False, False,  True, False, False], dtype=bool)

# here y+z==2 returns a boolean array that can be used to index x
>>> x[y + z == 2]
array([3, 7])

Upvotes: 4

Related Questions