CODEWITHSUNDEEP
pythonnumpy
user10652346
user10652346

Reputation:

Extract values that satisfy a condition from numpy array

Say I have the following arrays:

a = np.array([1,1,1,2,2,2])
b = np.array([4,6,1,8,2,1])

Is it possible to do the following:

a[np.where(b>3)[0]]
#array([1, 1, 2])

Thus select values from a according to the indices in which a condition in b is satisfied, but using exclusively np.where or a similar numpy function?

In other words, can np.where be used specifying only an array from which to get values when the condition is True? Or is there another numpy function to do this in one step?

Upvotes: 3

Views: 18121

Answers (3)

MB-F
MB-F

Reputation: 23637

Yes, there is a function: numpy.extract(condition, array) returns all values from array that satifsy the condition.

There is not much benefit in using this function over np.where or boolean indexing. All of these approaches create a temporary boolean array that stores the result of b>3. np.where creates an additional index array, while a[b>3]and np.extract use the boolean array directly.

Personally, I would use a[b>3] because it is the tersest form.

Upvotes: 8

cph_sto
cph_sto

Reputation: 7585

Let's use list comprehension to solve this -

a = np.array([1,1,1,2,2,2])
b = np.array([4,6,1,8,2,1])
indices = [i for i in range(len(b)) if b[i]>3]  # Returns indexes of b where b > 3 - [0, 1, 3]
a[indices]
    array([1, 1, 2])

Upvotes: 1

timgeb
timgeb

Reputation: 78650

Just use boolean indexing.

>>> a = np.array([1,1,1,2,2,2])                                                                                                   
>>> b = np.array([4,6,1,8,2,1])                                                                                                   
>>>                                                                                                                               
>>> a[b > 3]                                                                                                                      
array([1, 1, 2])

b > 3 will give you array([True, True, False, True, False, False]) and with a[b > 3] you select all elements from a where the indexing array is True.

Upvotes: 5

Related Questions