python - 将 ndarray 与一维数组中的值进行比较以获得掩码

Question

我有两个 numpy 数组，分别是 2D 和 1D。我想获得一个 2D 二进制掩码，其中掩码的每个元素如果与 1D 数组的任何元素匹配，则为 true。

示例

 2D array 
 -----------
 1 2 3 
 4 9 6
 7 2 3


 1D array 
 -----------
 1,9,3

 Expected output
 ---------------
 True   False  True
 False  True   False
 False  False  True

谢谢

Answer 1

您可以使用np.in1d。虽然 np.in1d 返回一个一维数组，但您可以随后简单地重新调整结果:

In [174]: arr = np.array([[1,2,3],[4,9,6],[7,2,3]])

In [175]: bag = [1,9,3]

In [177]: np.in1d(arr, bag).reshape(arr.shape)
Out[177]: 
array([[ True, False,  True],
       [False,  True, False],
       [False, False,  True]], dtype=bool)

请注意，in1d 正在检查 arr 中的元素是否与 bag 中的任何元素匹配。相反，arr == bag 测试 arr 的元素是否等于 bag 的广播元素逐元素 。您可以通过排列 bag 来查看差异:

In [179]: arr == np.array([1,3,9])
Out[179]: 
array([[ True, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)

In [180]: np.in1d(arr, [1,3,9]).reshape(arr.shape)
Out[180]: 
array([[ True, False,  True],
       [False,  True, False],
       [False, False,  True]], dtype=bool)

<小时/>

当您比较两个形状不等的数组时，NumPy 会尝试 broadcast在测试相等性之前将两个数组转换为单个兼容形状。在本例中，[1, 3, 9] 被广播到

array([[1, 3, 9],
       [1, 3, 9],
       [1, 3, 9]])

因为在左侧添加了新轴。可以这样查看广播效果:

In [181]: np.broadcast_arrays(arr, [1,3,9])
Out[185]: 
[array([[1, 2, 3],
        [4, 9, 6],
        [7, 2, 3]]), 
 array([[1, 3, 9],
        [1, 3, 9],
        [1, 3, 9]])]

一旦两个数组被广播到一个共同的形状，就测试相等性逐元素，这意味着测试相应位置的值平等。例如，在顶行中，相等测试为 1 == 1、2 == 3、3 == 9。因此，

In [179]: arr == np.array([1,3,9])
Out[179]: 
array([[ True, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)