布尔 2D numpy 数组的掩码，其中包含另一个 1D numpy 数组中的元素的值为 True答案

【问题标题】：Mask of boolean 2D numpy array with True values for elements contained in another 1D numpy array布尔 2D numpy 数组的掩码，其中包含另一个 1D numpy 数组中的元素的值为 True
【发布时间】：2022-01-18 19:43:46
【问题描述】：

举个例子。我有一个数组test，并且想要为所有等于ref 的元素的元素获取一个带有True 的布尔掩码。

import numpy as np
test = np.array([[2, 3, 1, 0], [5, 4, 2, 3], [6, 7, 5 ,4]])
ref = np.array([3, 4, 5])

我正在寻找类似的东西

mask = (test == ref[0]) | (test == ref[1]) | (test == ref[2])

在这种情况下应该产生

>>> print(mask)
[[False,  True, False, False],
 [ True,  True, False,  True],
 [False, False,  True,  True]]

但不必求助于任何循环。

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

Numpy 有一个函数 isin 可以做到这一点

np.isin(test, ref)

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

【讨论】：

np.in1d（由isin使用）如果ref与test相比足够小，则重复mask |= (ar1 == a)。

【解决方案2】：

你可以使用 numpy 广播：

mask = (test[:,None] == ref[:,None]).any(1)

输出：

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

注意。这比numpy.isin 更快，但是会创建一个 (X, X, Y) 大小的中间数组，其中 X, Y 是测试的形状，所以这会在非常大的数组上消耗一些内存

【讨论】：