Check if each element in a numpy array is in a separate list

Question

I'd like to do something like this:

>>> y = np.arange(5)
>>> y in (0, 1, 2)
array([True, True, True, False, False])

This syntax doesn't work. What's the best way to achieve the desired result?

(I'm looking for a general solution. Obviously in this specific case I could do y < 3.)

I'll spell this out a little more clearly for you guys, since at least a few people seem to be confused.

Here is a long way of getting my desired behavior:

new_y = np.empty_like(y)
for i in range(len(y)):
    if y[i] in (0, 1, 2):
        new_y[i] = True
    else:
        new_y[i] = False

I'm looking for this behavior in a more compact form.

Here's another solution:

new_y = np.array([True if item in (0, 1, 2) else False for item in y])

Again, just looking for a simpler way.

np.in1d(y, [0,1,2])

user3483203
– user3483203

2019-03-09 20:00:19 +00:00
Commented Mar 9, 2019 at 20:00 — user3483203
– user3483203, Commented Mar 9, 2019 at 20:00

beyarkay · Accepted Answer · 2023-05-08 04:46:44Z

2

A good general purpose tool is a broadcasted, or 'outer', comparison between elements of two arrays:

In [35]: y=np.arange(5)                                                         
In [36]: x=np.array([0,1,2])                                                    
In [37]: y[:,None]==x                                                           
Out[37]: 
array([[ True, False, False],
       [False,  True, False],
       [False, False,  True],
       [False, False, False],
       [False, False, False]])

This is doing a fast comparison between every element of y and every element of x. Depending on your needs, you can condense this array along one of the axes:

In [38]: (y[:,None]==x).any(axis=1)                                             
Out[38]: array([ True,  True,  True, False, False])

A comment suggested in1d. I think it's a good idea to look at its code. It has several strategies depending on the relative sizes of the inputs.

In [40]: np.in1d(y,x)                                                           
Out[40]: array([ True,  True,  True, False, False])
In [41]: np.array([True if item in x else False for item in y])                 
Out[41]: array([ True,  True,  True, False, False])

Which is fastest may depend on the size of the inputs. Starting lists your list comprehension might be faster. This pure list version is by far the fastest:

[True if item in (0,1,2) else False for item in (0,1,2,3,4)] 
[item in (0,1,2) for item in (0,1,2,3,4)]    # simpler

edited May 8, 2023 at 4:46

beyarkay

1,1631 gold badge10 silver badges28 bronze badges

answered Mar 9, 2019 at 20:41

hpaulj

233k14 gold badges260 silver badges392 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Starman Over a year ago

If anyone else ends up here later, and sees the np.in1d comment, probably better to check out np.isin as the modern successor: numpy.org/doc/stable/reference/generated/…

hpaulj Over a year ago

@Starman, but look at its code

beyarkay Over a year ago

I looked at the code, and np.isin just seems to call np.in1d under the hood? github.com/numpy/numpy/blob/v1.24.0/numpy/lib/… Also the np.in1d docs recommend using np.isin.

Collectives™ on Stack Overflow

Check if each element in a numpy array is in a separate list

1 Answer 1

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related