Использование np.where для поиска элемента из подмассива

Question

Я пытаюсь найти все элементы в массиве X, которые соответствуют элементу в массиве Y, используя np.where () и условие для функции where () сравнивает список (a), а не один элемент. Пожалуйста, смотрите следующий код:

X = np.array([[0, 2], [2, 1], [1, 3], [5, 9], [6, 7], [4, 6]])
Y = np.array([1, 2, 3, 4, 4, 5])
a = [2, 3, 4]
matchedX = X[np.where(Y == a)]

Я ожидаю получить такой результат:

array([[2, 1],
   [1, 3],
   [5, 9],
   [6, 7]])

но я получил другие результаты:

array([], shape=(0, 2), dtype=int64)

Итак, мне нужно альтернативное решение, чтобы я мог получить те же элементы, если я не знаю о значениях? Эта строка ниже дает мне точные результаты, которые я хочу, но я не знаю в предыдущих значениях a.

matchedX = X[np.where((Y == 2) | (Y==3) | (Y==4))]

Answer 1

Вы можете пропустить np.where, который здесь избыточен, и просто проиндексировать, используя np.isin:

X[np.isin(Y,a)]

array([[2, 1],
       [1, 3],
       [5, 9],
       [6, 7]])

Это потому, что np.isin дает вам логический массивгде Y находится в a:

array([False,  True,  True,  True,  True, False])

Таким образом, индексируя с помощью этого массива, он выбирает только те строки, где True

Answer 2

Вы можете использовать установленные функции numpy:

X[np.where(np.isin(Y, a))]

array([[2, 1],
       [1, 3],
       [5, 9],
       [6, 7]])