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

Question

我正试图使用np.where((和条件where((函数来查找数组X中与数组Y上的元素匹配的所有元素，该函数比较的是列表(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:的集合函数

X[np.where(np.isin(Y, a))]
array([[2, 1],
[1, 3],
[5, 9],
[6, 7]])

Answer 2

您可以跳过这里多余的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所在的行