如何快速检查2D numpy数组的每行是否存在于另一个2D数组中

问题原因说明 #

你之前使用np.isin得到逐元素匹配结果，是因为该方法默认判断单个元素是否存在于另一个数组的所有元素集合中，不会按行整体匹配。

推荐实现方法 #

方法1：广播匹配（小数据量场景，代码简洁） #

利用Numpy广播机制对两行全元素匹配，无需遍历：

import numpy as np

a = np.array([[1,0], [2,0], [3,0]])
b = np.array([[1,0], [3,0]])

result = (a[:, np.newaxis] == b).all(axis=2).any(axis=1)
print(result)
# 输出：array([ True, False,  True])

逻辑说明：

• a[:, np.newaxis] 将a的维度从(N, D)扩展为(N, 1, D)，和维度为(M, D)的b广播后得到(N, M, D)的比较矩阵
• all(axis=2) 按最后一个维度（行元素维度）判断是否全相等，得到(N, M)的匹配矩阵，每个元素代表a的某行是否和b的某行完全一致
• any(axis=1) 按a的行维度判断是否至少有一个匹配，得到最终的(N,)布尔数组

方法2：结构化数组视图（大数据量场景，性能更优） #

如果数组规模很大，广播会产生较大的中间矩阵，此时可以将每行转为单个结构化元素，直接用np.isin实现整行匹配，时间和空间复杂度更低：

def row_in_array(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    # 检查两个数组的列数是否一致
    assert a.shape[1] == b.shape[1], "两个数组的列数必须相等"
    # 将每行转为单个void类型的结构化元素
    dtype = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
    a_view = a.view(dtype).ravel()
    b_view = b.view(dtype).ravel()
    return np.isin(a_view, b_view)

# 测试
result = row_in_array(a, b)
print(result)
# 输出：array([ True, False,  True])

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