Extract values that satisfy a condition from numpy array

Admin 1
Viewed 10251

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?

numpy python
3 Answers

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.

user_30051670

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.

user_36200030

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])
user_83292130
Related