CODEWITHSUNDEEP
pythonnumpy
Hammer
Hammer

Reputation: 10319

options for applying conditions to numpy arrays

I want to apply conditions to a numpy array and I feel like there is a better way out there. As a toy example say I want to know where the elements are equal to 2 or 3.

import numpy as np
a = np.arange(5)

one way would be to construct my condition piece by piece with numpy functions like so

result = np.logical_or(a == 2, a == 3)

One can see how this could get unwieldy with more complicated conditions though. Another option would be to use list comprehensions

result = np.array([x for x in a if x == 2 or x==3])

which is nice because now all my conditional logic can live together in one place but feels a little clunky because of the conversion to and from a list. It also doesn't work too well for multidimensional arrays.

Is there a better alternative that I am missing?

Upvotes: 1

Views: 7624

Answers (3)

dawg
dawg

Reputation: 103694

You can & together views to get arbitrarily complex results:

>>> A = np.random.randint(0, 100, 25).reshape(5,5)
>>> A
array([[98,  4, 46, 40, 24],
       [93, 75, 36, 19, 63],
       [23, 10, 62, 14, 59],
       [99, 24, 57, 78, 74],
       [ 1, 83, 52, 54, 27]])
>>> A>10
array([[ True, False,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [ True, False,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [False,  True,  True,  True,  True]], dtype=bool)
>>> (A>10) & (A<20)
array([[False, False, False, False, False],
       [False, False, False,  True, False],
       [False, False, False,  True, False],
       [False, False, False, False, False],
       [False, False, False, False, False]], dtype=bool)
>>> (A==19) | (A==14)  # same output

You can also write a function and use map to call the function on each element. Inside the function have as many tests as you wish:

>>> def test(ele):
...    return ele==2 or ele==3
... 
>>> map(test,np.arange(5))
[False, False, True, True, False]

You can use numpy.vectorize:

>>> def test(x):
...    return x>10 and x<20
...
>>> v=np.vectorize(test)
>>> v(A)
array([[False, False, False, False, False],
       [False, False, False,  True, False],
       [False, False, False,  True, False],
       [False, False, False, False, False],
       [False, False, False, False, False]], dtype=bool)

Upvotes: 1

kiriloff
kiriloff

Reputation: 26335

you can remove elements form numpy array with delete

np.delete(a,[0,1,4])

or if you want to keep with the complement,

np.delete(a,np.delete(a,[2,3]))

Upvotes: 1

mgilson
mgilson

Reputation: 309821

It's useful to point out that in the first example, you have a logical array, not the array [2, 3] (like you get in the second example). To recover the result from the second answer, you'd need 

result = a[result]

However, in this case, since you're using boolean masks (True/False approximately equivalent to 1/0), you can actually use bitwise or to do the same thing as logical_or:

result = a[(a==2) | (a==3)]

A word of caution here -- Make sure you use parenthesis. Otherwise, operator precedence can be a bit nasty in these expressions (| binds tighter than ==).

Upvotes: 3

Related Questions