如何高效判断2D numpy数组的每个元素是否存在于指定集合中

直接使用Numpy原生提供的np.isin()方法即可实现需求，该方法是向量化实现的，性能远高于自行遍历数组，返回结果的维度和输入数组完全一致。

示例代码 #

和你给出的需求完全匹配的实现如下：

import numpy as np

segmask = np.array([[1,2,3,4]])
numbers = {2,4}

# 你需要的check_if_in_set函数可以直接封装为：
def check_if_in_set(target_set, arr):
    return np.isin(arr, target_set)

# 调用后得到结果
result = check_if_in_set(numbers, segmask)

输出结果 #

>>> print(result)
[[False  True False  True]]

实用提示 #

• np.isin支持直接接收集合作为第二个参数，不需要额外转成列表
• 如果需要判断元素不在目标集合中，直接对结果取反即可：result = ~check_if_in_set(numbers, segmask)
• 处理无重复元素的超大数组时，可以添加assume_unique=True参数进一步提升运行速度：np.isin(arr, target_set, assume_unique=True)

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