高效筛选两个二维NumPy数组中在A不在B的行

针对你提到的大型二维数组（如(50000, 2)），嵌套循环的效率瓶颈在于Python层面的循环开销，改用NumPy的向量化操作能大幅提升性能，以下是几种实用的高效实现方法：

方法1：利用view+setdiff1d（推荐） #

该方法通过将每行转换为不可变的字节视图，把二维行的比较转化为一维元素的集合差运算，完全利用NumPy的底层优化：

import numpy as np

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

# 将二维数组的每行转为单个字节类型的元素（保持行的唯一性）
row_dtype = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
a_rows = a.view(row_dtype).ravel()
b_rows = b.view(row_dtype).ravel()

# 找出a中独有的行标识
unique_a_rows = np.setdiff1d(a_rows, b_rows, assume_unique=False)

# 根据标识筛选原数组中的行
in_a_and_not_b = a[np.in1d(a_rows, unique_a_rows)]
print(in_a_and_not_b)
# 输出：[[3 4]
#        [7 8]]

原理说明： #

• view操作不复制数据，只是改变数组的内存解读方式，把每行的所有字节打包成一个void类型的元素，这样每行就可以被当成单个值处理。
• np.setdiff1d快速计算两个一维数组的集合差，得到仅在a_rows中存在的元素。
• np.in1d匹配原数组中对应的行，完成最终筛选。

方法2：结合np.unique与布尔索引 #

如果数组a中存在重复行，且需要保留所有重复项，可先提取a的唯一行，再筛选不在b中的行，最后找回原数组中的对应行：

# 获取a的唯一行及对应索引
a_unique, idx = np.unique(a, axis=0, return_inverse=True)

# 检查唯一行是否不在b中
mask = ~np.isin(a_unique.view(row_dtype), b_rows)

# 筛选出符合条件的行并还原原数组中的重复项
in_a_and_not_b = a_unique[mask][idx[np.in1d(idx, np.where(mask)[0])]]

性能对比 #

对于(50000, 2)规模的数组，上述方法的耗时通常在毫秒级，而嵌套循环可能需要数十秒甚至更久，核心原因是NumPy的向量化操作直接在C层面执行，避免了Python循环的额外开销。

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