python - 使用 np.where 从子数组中查找元素

Question

我正在尝试使用 np.where() 查找数组 X 中与数组 Y 中的元素匹配的所有元素，并且 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 的值，我可以获得相同的元素？下面这一行给出了我想要的确切结果，但我不知道以前的 a 值。

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

Answer 1

可以使用numpy的set函数:

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

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