numpy条件列表成员资格元素明智答案

【问题标题】：numpy conditional list membership element wisenumpy条件列表成员资格元素明智
【发布时间】：2018-06-27 19:37:07
【问题描述】：

我有一个 2D numpy 数组：

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

我有一个值列表要保留：

vals_keep = [1,2]

我想测试数组中每个元素的列表成员资格。比如：

mask = a in vals_keep

我想要的结果：

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

【问题讨论】：

标签： python numpy numpy-ndarray

【解决方案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]]]])

【讨论】：

【解决方案2】：

这是使用广播的一种方式：

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

对于小型数组（少于 100 行），这比 isin 快：

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 广播对于大型数组/矩阵来说效率不高。

【讨论】：