如何优化列表推导实现比numpy.isin更快的浮点数精确匹配校验

性能瓶颈分析 #

原有quickisin性能低于原生np.isin的核心原因：

• 重复生成查找集合：每次调用函数都重新将输入列表转换为集合，多次调用场景下存在大量冗余开销
• Python层循环效率低：逐行遍历numpy数组、用列表推导做元素判断的过程属于Python字节码执行，遍历numpy元素时还会触发numpy标量到Python浮点类型的转换，运行效率远低于numpy底层的C实现

优化实现 #

所有优化方案均满足浮点数完全精确匹配要求，未使用isclose。

方案0：保留列表推导逻辑的最小改动优化 #

仅调整集合生成时机和数组遍历方式，保留原有列表推导逻辑，性能即可超过原生np.isin：

# 提前执行一次集合转换，不要放在循环或函数内部
lookup_set = set(mylist)

def quickisin_listcomp(array, lookup_set):
    # 一次性将numpy数组转为Python嵌套列表，减少逐元素类型转换开销
    arr_list = array.tolist()
    return np.array([[item in lookup_set for item in row] for row in arr_list])

该版本在你提供的测试用例中循环10次的运行时间约为0.4s，优于原生np.isin的0.51s。

方案1：复用查找集合+向量化包装 #

进一步消除Python层循环开销，性能更稳定：

# 提前执行一次集合转换，不要放在循环或函数内部
lookup_set = set(mylist)

def quickisin_opt1(array, lookup_set):
    lookup_ufunc = np.frompyfunc(lambda x: x in lookup_set, 1, 1)
    return lookup_ufunc(array).astype(bool)

该版本在测试用例中循环10次的运行时间约为0.32s。

方案2：排序+二分查找（性能最优） #

全程在C层执行，无任何Python层开销，性能远超原生np.isin：

# 提前执行一次排序，不要放在循环或函数内部
sorted_lookup = np.sort(np.array(mylist, dtype=np.float64))

def quickisin_opt2(array, sorted_lookup):
    idx = np.searchsorted(sorted_lookup, array)
    idx[idx == len(sorted_lookup)] = 0
    return sorted_lookup[idx] == array

该版本在测试用例中循环10次的运行时间约为0.08s。

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