python - numpy 条件列表成员元素明智

我有一个 2D numpy 数组:

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

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

vals_keep = [1,2]

我想测试数组中每个元素的列表成员资格。像这样的东西:

mask = a in vals_keep

我想要的结果:

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

最佳答案

您可以使用 isin

isin is an element-wise function version of the python keyword 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]]]])

关于python - numpy 条件列表成员元素明智，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/51070234/