numpy条件列表会员元素明智

Question

我有一个2D numpy数组:

a = np.array([[0,1],
              [2,3]])

我有一个要保留的值列表:

vals_keep = [1,2]

我想测试数组中每个元素的列表成员资格.就像是:

mask = a in vals_keep

我想要的结果:

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

Answer 1

您可以使用 isin

isin 是python关键字的元素功能版本 in

np.isin(a, vals_keep)

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

另一个好处isin是它可以灵活地使用不同尺寸的数组:

a = np.arange(4).reshape(1,2,2,1)
np.isin(a, vals_keep)

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

        [[ True],
         [False]]]])

Answer 2

以下是使用广播的一种方式:

In [35]: (a[:, :, None] == vals_keep).any(2)
Out[35]: 
array([[False,  True],
       [ True, False]])

哪个比isin小数组(少于100行)快:

In [37]: %timeit np.isin(a, vals_keep)
22 µs ± 728 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [38]: %timeit (a[:, :, None] == vals_keep).any(2)
12.6 µs ± 95.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

对于大型阵列,最好使用,isin因为3D中的广播对于大型阵列/矩阵来说效率不高.