How do I select elements of an array given condition?

I like to use np.vectorize for such tasks. Consider the following:

>>> # Arrays
>>> x = np.array([5, 2, 3, 1, 4, 5])
>>> y = np.array(['f','o','o','b','a','r'])

>>> # Function containing the constraints
>>> func = np.vectorize(lambda t: t>1 and t<5)

>>> # Call function on x
>>> y[func(x)]
>>> array(['o', 'o', 'a'], dtype='<U1')

The advantage is you can add many more types of constraints in the vectorized function.

Hope it helps.


Add one detail to @J.F. Sebastian's and @Mark Mikofski's answers:
If one wants to get the corresponding indices (rather than the actual values of array), the following code will do:

For satisfying multiple (all) conditions:

select_indices = np.where( np.logical_and( x > 1, x < 5) )[0] #   1 < x <5

For satisfying multiple (or) conditions:

select_indices = np.where( np.logical_or( x < 1, x > 5 ) )[0] # x <1 or x >5

IMO OP does not actually want np.bitwise_and() (aka &) but actually wants np.logical_and() because they are comparing logical values such as True and False - see this SO post on logical vs. bitwise to see the difference.

>>> x = array([5, 2, 3, 1, 4, 5])
>>> y = array(['f','o','o','b','a','r'])
>>> output = y[np.logical_and(x > 1, x < 5)] # desired output is ['o','o','a']
>>> output
array(['o', 'o', 'a'],
      dtype='|S1')

And equivalent way to do this is with np.all() by setting the axis argument appropriately.

>>> output = y[np.all([x > 1, x < 5], axis=0)] # desired output is ['o','o','a']
>>> output
array(['o', 'o', 'a'],
      dtype='|S1')

by the numbers:

>>> %timeit (a < b) & (b < c)
The slowest run took 32.97 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 1.15 µs per loop

>>> %timeit np.logical_and(a < b, b < c)
The slowest run took 32.59 times longer than the fastest. This could mean that an intermediate result is being cached.
1000000 loops, best of 3: 1.17 µs per loop

>>> %timeit np.all([a < b, b < c], 0)
The slowest run took 67.47 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 5.06 µs per loop

so using np.all() is slower, but & and logical_and are about the same.


Your expression works if you add parentheses:

>>> y[(1 < x) & (x < 5)]
array(['o', 'o', 'a'], 
      dtype='|S1')

Tags:

Python

Numpy