如何高效实现NumPy数组与标签列表对比以生成布尔掩码？

直接用NumPy的np.isin()函数就能解决问题，它完全替代循环逻辑，利用矢量化操作大幅提升处理速度：

import numpy as np
input_array = np.arange(15, dtype=np.int64).reshape(3, 5)
good_labels = [1, 5, 7]
mask_output = np.isin(input_array, good_labels)

原循环方法慢的原因 #

你之前的循环会对每个标签单独遍历整个input_array做比较，当标签列表规模较大时，重复遍历的开销会急剧上升。而np.isin()是底层矢量化实现，一次性完成所有元素的判断，效率远高于Python循环。

为什么input_array in good_labels不生效？ #

这是Python原生的成员判断逻辑，它只会检查整个数组对象是否在good_labels里，不会逐元素判断，所以返回的是单个布尔值，不是你需要的同形状布尔数组。

进阶优化 #

如果good_labels是超大列表，先转成集合再传入np.isin()，能进一步提升查找效率：

good_labels_set = set(good_labels)
mask_output = np.isin(input_array, good_labels_set)

内容的提问来源于stack exchange，提问作者toto