Учитывая два ndarrays, как я могу извлечь записи, которые принадлежат к подмножеству индексов способом numpy-ish?

Question

С учетом двух ndarrays (10000, 2000) и еще одного (10000,)

Второй имеет 10 различных классов ("class1", "class2", ...), смешанных вместе (без порядка). Индексы второго массива соответствуют первому затемнению первого массива.

Как мне извлечь индексы ТОЛЬКО class1 и class2 и выбрать подходящие записи из массива 1 способом numpy-ish?

например, для

["class1", "class3", "class4", "class2", "class1"]

и

[
[1,1,1,1]
[2,2,2,2]
[3,3,3,3]
[4,4,4,4]
[5,5,5,5]
]

Я бы получил вывод

[
[1,1,1,1]
[4,4,4,4]
[5,5,5,5]
]

Answer 1

Вы можете использовать Логическое индексирование через np.isin:

import numpy as np

classes = np.array(["class1", "class3", "class4", "class2", "class1"])
data = np.array([[1,1,1,1], [2,2,2,2], [3,3,3,3],
                 [4,4,4,4], [5,5,5,5]])

bool_idx = np.isin(classes, ['class1', 'class2'])
res = data[bool_idx]

# array([[1, 1, 1, 1],
#        [4, 4, 4, 4],
#        [5, 5, 5, 5]])