如何高效实现多维数组场景下替代numpy.isin的布尔掩码生成？

可以利用NumPy的广播机制实现无列表推导的高效运算，核心思路是通过维度扩展让b的每个子数组与a的所有元素进行逐元素比较，再聚合结果得到对应掩码行。

实现代码 #

import numpy as np

a = np.array([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
b = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]])

# 通过广播实现逐组匹配
mask = (b[:, :, None] == a.T).any(axis=1)
print(mask)

原理说明 #

1. 维度扩展：b[:, :, None]将b从(5,2)扩展为(5,2,1)，a.T将a从(1,10)转置为(10,1)；
2. 广播比较：两者进行相等比较时，广播机制会自动扩展维度，得到形状为(5,2,10)的布尔数组，其中每个b的元素都会与a的所有元素逐一对比；
3. 聚合结果：沿axis=1（即每个子数组的维度）调用any()，只要子数组中有一个元素匹配a的对应位置，就标记为True，最终得到形状为(5,10)的目标掩码。

性能优势 #

广播是NumPy底层优化的运算方式，避免了Python层面的循环（包括列表推导），处理大型数组时效率远高于纯Python循环实现。

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