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

时间:2018-09-18 10:09:21

标签: python numpy

让我们说我们有两个整数numpy数组A和B,大小分别为(N,M)。我想检查N中的每个i是否为A[i,:]中的B[i,:]

一个for循环的实现是:

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

但是,这对于大型阵列来说相当慢。有没有更快的方法来实现这一目标? (例如矢量化?)

谢谢!

3 个答案:

答案 0 :(得分:2)

这是一种基于逐行偏移的矢量化方法,如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]])

答案 1 :(得分: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“压缩”在元组中并对其进行操作,而不会产生这样做的开销。因此,我们检查每个项目及其行索引压缩后是否包含在一维集中;因此只有在行索引和数值一致时才匹配。

Divakars解决方案可能具有速度优势,但是两种解决方案应具有相同的时间复杂性。并且这里发布的解决方案适用于任意dtypes。甚至您尝试匹配的代币本身都是ndarrays而不是标量!

答案 2 :(得分:0)

也许是这样的:

>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])