匹配两个2D数组的行并使用numpy获取行索引图

发布于 2021-01-29 14:59:26

假设您有两个2D数组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检查2d
numpy数组中的成员资格
)检查A的行是否在B中,但这不会产生索引图。

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

一种替代方法是使用pandas merge函数,但我只想使用numpy做到这一点。

关注者
0
被浏览
119
1 个回答
  • 面试哥
    面试哥 2021-01-29
    为面试而生,有面试问题,就找面试哥。

    方法1

    这是基于的一个views。利用np.argwheredocs)返回满足条件的元素的索引,在这种情况下为成员资格。--

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


知识点
面圈网VIP题库

面圈网VIP题库全新上线,海量真题题库资源。 90大类考试,超10万份考试真题开放下载啦

去下载看看