python - 检查矩阵的每一行是否在另一个矩阵中 [Python]

Question

我有两个矩阵(或者更好，2D numpy 数组)，A 和 B，具有相同的列数(好吧，公平地说，它们有不同的列)行数，但是我只是将两个矩阵切成两列)但行数不同。

我想检查A中的每一行是否都在B中，以及B中正确的行索引是什么。我实际上使用的是常规循环，即

for r in range(A.shape[0]):
    idx=numpy.where(numpy.all(B[:,[0,1]]==A[r,[1,2]],axis=1))
    idx=idx[0][0]

我想知道是否有一些更智能的技术可以做到这一点，可以使用索引方法并避免循环。我知道 numpy 的 in1d，但它只适用于一维数组。

对于熟悉 Matlab 的读者，我正在寻找触发 'rows' 选项时的 ismember() 行为。

Answer 1

您可以使用 NumPy broadcasting 来使用矢量化方法，就像这样 -

np.argwhere((B[:,None,[0,1]] == A[:,[1,2]]).all(-1))

或者，由于您正在处理仅包含 2 元素的行，因此可以建议采用保持 2D 的内存高效方法，如下所示 -

np.argwhere((B[:,None,0] == A[:,1]) & (B[:,None,1] == A[:,2]))

输出将是 (N,2) 形状的数组，其中第一列为我们提供 B 的行索引，第二列为 A 对应于一行中所有元素的匹配。

示例运行 -

In [154]: A[:,[1,2]]
Out[154]: 
array([[0, 1],
       [0, 2],
       [2, 1],
       [1, 2],
       [0, 1]])

In [155]: B[:,[0,1]]
Out[155]: 
array([[0, 1],
       [2, 2],
       [1, 2],
       [0, 2],
       [2, 1],
       [2, 1],
       [1, 0]])

In [156]: np.argwhere((B[:,None,[0,1]] == A[:,[1,2]]).all(-1))
Out[156]: 
array([[0, 0],
       [0, 4],
       [2, 3],
       [3, 1],
       [4, 2],
       [5, 2]])