如何通过与一维列表对比实现二维NumPy数组的掩码操作？

需求说明 #

给定二维NumPy数组：

[1 2 3 1]
[2 3 1 2]
[3 1 2 3]
[1 2 3 1]

和一维整数列表[1, 3]，需要生成一个布尔数组，原数组中属于该列表的元素对应True，期望输出：

[True False True True]
[False True True False]
[True True False True]
[True False True True]

现有问题 #

当前采用循环遍历列表结合np.logical_or的方法，代码如下：

boolean_array = array == mask_list[0]
for list_element in mask_list[0:]:
    boolean_array = np.logical_or(boolean_array, array == list_element)

return boolean_array

但当二维数组是大图像、列表元素可达300个时，运行速度极慢，需要找到替代循环的高效NumPy操作。

高效解决方案 #

直接使用NumPy内置的np.isin()函数，这是专门用于判断数组元素是否存在于目标列表的矢量化操作，完全替代循环，性能提升显著。

代码示例 #

import numpy as np

# 定义输入数组和掩码列表
array = np.array([
    [1, 2, 3, 1],
    [2, 3, 1, 2],
    [3, 1, 2, 3],
    [1, 2, 3, 1]
])
mask_list = [1, 3]

# 生成布尔掩码数组
boolean_array = np.isin(array, mask_list)

print(boolean_array)

输出结果 #

[[ True False  True  True]
 [False  True  True False]
 [ True  True False  True]
 [ True False  True  True]]

原理说明 #

np.isin()是矢量化操作，底层由C语言实现，无需Python层面的循环遍历，能一次性完成整个数组的元素判断。无论你的数组是百万像素的图像，还是包含数百元素的掩码列表，都能以最优效率执行，比循环+np.logical_or的方式快几个数量级。

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