python之匹配两个二维数组的行并使用 numpy 获取行索引映射

假设您有两个二维数组 A 和 B，并且您想要检查 B 中包含一行 A 的位置。 如何使用 numpy 最有效地做到这一点？

例如

a = np.array([[1,2,3], 
              [4,5,6], 
              [9,10,11]]) 
 
b = np.array([[4,5,6], 
              [4,3,2], 
              [1,2,3], 
              [4,8,9]]) 
map = [[0,2], [1,0]]  # row 0 of a is at row index 2 of array B 

我知道如何使用 in1d ( test for membership in a 2d numpy array ) 检查 A 的一行是否在 B 中，但这不会产生索引映射。

此映射的目的是(最终)根据某些列将两个数组合并在一起。
当然可以逐行执行此操作，但这会变得非常低效，因为我的数组具有 (50 Mio., 20) 的形状。

另一种方法是使用 pandas merge function ，但我只想使用 numpy 来执行此操作。

请您参考如下方法:

方法 #1

这是一个基于views 的。使用 np.argwhere ( docs ) 返回满足条件的元素的索引，在本例中为成员资格。 -

def view1D(a, b): # a, b are arrays 
    a = np.ascontiguousarray(a) 
    b = np.ascontiguousarray(b) 
    void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1])) 
    return a.view(void_dt).ravel(),  b.view(void_dt).ravel() 
 
def argwhere_nd(a,b): 
    A,B = view1D(a,b) 
    return np.argwhere(A[:,None] == B) 

方法 #2

这是另一个O(n)，因此性能要好得多，尤其是在大型阵列上 -

def argwhere_nd_searchsorted(a,b): 
    A,B = view1D(a,b) 
    sidxB = B.argsort() 
    mask = np.isin(A,B) 
    cm = A[mask] 
    idx0 = np.flatnonzero(mask) 
    idx1 = sidxB[np.searchsorted(B,cm, sorter=sidxB)] 
    return idx0, idx1 # idx0 : indices in A, idx1 : indices in B 

方法 #3

另一个 O(n) 使用 argsort() -

def argwhere_nd_argsort(a,b): 
    A,B = view1D(a,b) 
    c = np.r_[A,B] 
    idx = np.argsort(c,kind='mergesort') 
    cs = c[idx] 
    m0 = cs[:-1] == cs[1:] 
    return idx[:-1][m0],idx[1:][m0]-len(A) 

样本运行与之前相同的输入 -

In [650]: argwhere_nd_searchsorted(a,b) 
Out[650]: (array([0, 1]), array([2, 0])) 
 
In [651]: argwhere_nd_argsort(a,b) 
Out[651]: (array([0, 1]), array([2, 0]))