Get indices of intersecting rows of Numpy 2d Array

前端 未结 2 830
攒了一身酷
攒了一身酷 2020-12-19 07:59

I want to get the indices of the intersecting rows of a main numpy 2d array A, with another one B.

A=array([[1, 2],
         [3, 4],
         [5, 6],
                


        
相关标签:
2条回答
  • 2020-12-19 08:23

    With minimal changes, you can get your approach to work:

    In [15]: A
    Out[15]: 
    array([[ 1,  2],
           [ 3,  4],
           [ 5,  6],
           [ 7,  8],
           [ 9, 10]])
    
    In [16]: B
    Out[16]: 
    array([[1, 4],
           [1, 2],
           [5, 6],
           [6, 3]])
    
    In [17]: np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1))
    Out[17]: array([ True, False,  True, False, False], dtype=bool)
    
    In [18]: np.nonzero(np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1)))
    Out[18]: (array([0, 2], dtype=int64),)
    
    In [19]: np.nonzero(np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1)))[0]
    Out[19]: array([0, 2], dtype=int64)
    

    If your arrays are not floats, and are both contiguous, then the following will be faster:

    In [21]: dt = np.dtype((np.void, A.dtype.itemsize * A.shape[1]))
    
    In [22]: np.nonzero(np.in1d(A.view(dt).reshape(-1), B.view(dt).reshape(-1)))[0]
    Out[22]: array([0, 2], dtype=int64)
    

    And a quick timing:

    In [24]: %timeit np.nonzero(np.in1d(A.view('i,i').reshape(-1), B.view('i,i').reshape(-1)))[0]
    10000 loops, best of 3: 75 µs per loop
    
    In [25]: %timeit np.nonzero(np.in1d(A.view(dt).reshape(-1), B.view(dt).reshape(-1)))[0]
    10000 loops, best of 3: 29.8 µs per loop
    
    0 讨论(0)
  • 2020-12-19 08:29

    You can use np.char.array() objects to do this comparison using np.in1d():

    s1 = np.char.array(A[:,0]) + '-' + np.char.array(A[:,1])
    s2 = np.char.array(B[:,0]) + '-' + np.char.array(B[:,1])
    
    np.where(np.in1d(s1, s2))[0]
    #array([0, 2], dtype=int64)
    

    NOTE: A and B must be of the same data type (int, float, etc) for this to work.

    0 讨论(0)
提交回复
热议问题