python - 根据值的集合成员资格为 numpy 数组创建掩码

Question

我想根据该数组的元素是否是某个集合的成员，为数组创建一个“掩码”索引数组。我想要的可以通过以下方式实现:

x = np.arange(20)
interesting_numbers = {1, 5, 7, 17, 18}
x_mask = np.array([xi in interesting_numbers for xi in x])

我想知道是否有更快的方法来执行最后一行。实际上，它通过重复调用 __contains__ 方法在 Python 中构建一个列表，然后将该列表转换为 numpy 数组。

我想要类似 x_mask = x[x ininteresting_numbers] 的内容，但这不是有效的语法。

Answer 1

您可以使用np.in1d:

np.in1d(x, list(interesting_numbers))
#array([False,  True, False, False, False,  True, False,  True, False,
#       False, False, False, False, False, False, False, False,  True,
#        True, False], dtype=bool)

时序，数组x越大，速度越快:

x = np.arange(10000)
interesting_numbers = {1, 5, 7, 17, 18}

%timeit np.in1d(x, list(interesting_numbers))
# 10000 loops, best of 3: 41.1 µs per loop

%timeit x_mask = np.array([xi in interesting_numbers for xi in x])
# 1000 loops, best of 3: 1.44 ms per loop