Numpy技巧：如何基于3D数组与1D数组创建N×M掩码

给定形状为(N,M,3)的3D numpy数组A，以及形状为(3,)的1D数组B，要生成一个N×M的布尔掩码，标记A中最后一维元素与B完全相等的位置，你可以用以下更简洁的实现：

方法一：利用广播与all()方法 #

import numpy as np

mask = (A == B).all(axis=2)

原理说明： #

• numpy的广播机制会自动将1D数组B扩展为与A同形状的(N,M,3)数组，逐元素完成相等性比较，得到一个同样形状的布尔数组。
• 调用all(axis=2)沿最后一维（即维度索引为2的轴）进行逻辑与判断，只要该维度上所有元素都为True，就会在对应位置生成True，最终得到(N,M)的掩码数组。

这种方法的优势在于扩展性强——哪怕最后一维的长度不是3，比如是5、10，代码都不需要修改，直接复用即可，完全避免了手动逐个维度写比较语句的繁琐。

补充：处理浮点数精度问题 #

如果你的数组是浮点数类型，直接用==比较可能会因为精度误差出现错误结果，这时可以用np.allclose()替代：

mask = np.allclose(A, B, atol=1e-8).all(axis=2)

atol参数可以根据你的精度需求调整，默认值是1e-08。

内容的提问来源于stack exchange，提问作者dima.rus