python - 有没有一种优雅的方法来检查数组中的哪些元素位于另一个容器中？

Question

我确信相同或类似的问题已经被问过很多次了，但也许我的措辞有问题，所以我还没有找到合适的帖子。

这是一个例子:

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

有没有比手动循环容器更好的方法？

Answer 1

使用np.isin :

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

输出

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

作为替代方案，扁平版本 np.in1d :

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

输出

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

更新

请注意，如果 test_elements(即 ['John', 'Laura'])是一个集合，则此方法将不起作用。来自文档:

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.