Is there an elegant way to check which elements in an array are in another container?

user_5472354 0
Viewed 39

I'm sure the same or a similar question has been asked many times, but maybe my phrasing is off, so I didn't find a fitting post yet.

Here is an example:

import numpy as np


students = np.array(
    [('Jim', 2), ('John', 3), ('Lisa', 5), ('Laura', 1)],
    dtype=[('name', 'O'), ('grade', 'i4')]
)

# this works
good_students = students[students['grade'] < 3]

# this doesn't because unfortunately, numpy doesn't understand that I want 
# it to do elementwise comparison
# absent_students = students[students['name'] in ['John', 'Laura']]

# this does but is ugly
absent_students = np.array(
    [], dtype=[('name', 'O'), ('grade', 'i4')]
)
for name in ['John', 'Laura']:
    absent_students = np.append(
        absent_students, students[students['name'] == name]
    )

Is there a better way than manually looping over the container?

numpy arrays python
1 Answers

Use np.isin:

absent_students = students[np.isin(students['name'],['John', 'Laura'])]
print(absent_students)

Output

[('John', 3) ('Laura', 1)]

As an alternative the flattened version, np.in1d:

absent_students = students[np.in1d(students['name'],['John', 'Laura'])]
print(absent_students)

Output

[('John', 3) ('Laura', 1)]

UPDATE

Notice that this won't work if the test_elements, i.e. ['John', 'Laura'], is a set. From the documentation:

If test_elements is a set (or other non-sequence collection) it will be converted to an object array with one element, rather than an array of the values contained in test_elements. This is a consequence of the array constructor’s way of handling non-sequence collections. Converting the set to a list usually gives the desired behavior.

user_40015920
Related