检查numpy数组是否在另一个numpy数组内,并创建一个掩码

时间:2019-09-30 15:05:53

标签: python numpy

给出numpy数组a

a = [[[0 0] [1 0] [2 0]]
     [[0 1] [1 1] [2 1]]
     [[0 2] [1 2] [2 2]]]

列表b

b = [[1, 0], [2, 0]]

我如何获得口罩c

c = [[False True True]
     [False False False]
     [False False False]]

1 个答案:

答案 0 :(得分:1)

您可以使用numpy广播功能将b中的每个数字对与b中的所有数字对进行比较

## np.newaxis add a new dimension at that position. missing dimension (i.e 
## dimension with size 1) will repeat to match size of corresponding dimension

a = np.asarray([[[0, 0], [1, 0], [2, 0]],
     [[0, 1], [1, 1], [2, 1]],
     [[0, 2], [1, 2], [2, 2]]])[:,:,np.newaxis,:]

b = np.array([[1, 0], [2, 0]])[np.newaxis,:,:]

(a == b).all(axis=3).any(axis=2)

结果

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