快速检查numpy数组的子维度中的元素是否在另一个numpy数组的子维度中的方法

Question

假设我们有两个整数numpy数组A和B大小（N，M）。我想检查N中的每一个，是A[i,:]中的B[i,:]。

for循环实现是：

for i in range(N):
    C[i] = np.isin(A[i,:],B[i,:])

然而，对于大型阵列来说这是非常慢的。有没有更快的方法来实现这个？ （例如矢量化？）

谢谢！

Answer 1

这是一种基于每行偏移的矢量化方法，如Vectorized searchsorted numpy's solution中更详细讨论的那样 -

# https://stackoverflow.com/a/40588862/ @Divakar
def searchsorted2d(a,b):
    m,n = a.shape
    max_num = np.maximum(a.max() - a.min(), b.max() - b.min()) + 1
    r = max_num*np.arange(a.shape[0])[:,None]
    p = np.searchsorted( (a+r).ravel(), (b+r).ravel() ).reshape(m,-1)
    return p - n*(np.arange(m)[:,None])

def numpy_isin2D(A,B):
    sB = np.sort(B,axis=1)
    idx = searchsorted2d(sB,A)
    idx[idx==sB.shape[1]] = 0
    return np.take_along_axis(sB, idx, axis=1) == A

样品运行 -

In [351]: A
Out[351]: 
array([[5, 0, 3, 3],
       [7, 3, 5, 2],
       [4, 7, 6, 8],
       [8, 1, 6, 7],
       [7, 8, 1, 5]])

In [352]: B
Out[352]: 
array([[8, 4, 3, 0, 3, 5],
       [0, 2, 3, 8, 1, 3],
       [3, 3, 7, 0, 1, 0],
       [4, 7, 3, 2, 7, 2],
       [0, 0, 4, 5, 5, 6]])

In [353]: numpy_isin2D(A,B)
Out[353]: 
array([[ True,  True,  True,  True],
       [False,  True, False,  True],
       [False,  True, False, False],
       [False, False, False,  True],
       [False, False, False,  True]])

Answer 2

numpy_indexed包（免责声明：我是它的作者）可用于获得类似于Divakars的解决方案，但抽象的低级细节：

import numpy_indexed as npi
Ar = np.indices(A.shape)[0]
Br = np.indices(B.shape)[0]
isin = npi.in_((A.flatten(), Ar.flatten()), (B.flatten(), Br.flatten())).reshape(A.shape)

numpy_indexed包中的所有函数在ndarrays上运行，或者在这种情况下，ndarrays的元组，实际上等同于“压缩”元组中的ndarrays并对其进行操作，而不会产生这样做的开销。因此，我们检查是否包含在1d集合中，每个项目用其行索引压缩;所以只有当行索引和数值都重合时才会记录匹配。

Divakars解决方案可能具有速度优势，但两种解决方案应该具有相同的时间复杂度。这里发布的解决方案适用于任意dtypes;或者即使你想要匹配的tokes是ndarrays本身而不是标量！

Answer 3

也许是这样的：

>A = np.zeros((5, 5))
>B = np.ones((5, 5))
>B[2, :] = 0
>print(A)
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]])
>print(B)
array([[1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.],
       [0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.]])
>(A == B).all(axis=1)
array([False, False,  True, False, False])