Find indexes of matching rows in two 2-D arrays

Question

Suppose that I have two 2-D arrays as follows:

array([[3, 3, 1, 0],
       [2, 3, 1, 3],
       [0, 2, 3, 1],
       [1, 0, 2, 3],
       [3, 1, 0, 2]], dtype=in         


        

Answer 1

You can use the void data type trick to use 1D functions on the rows of your two arrays. a_view and b_view are 1D vectors, each entry representing a full row. I have then chosen to sort an array and use np.searchsorted to find the items of the other array in that one. If the array we sort has length m and the other has length n, sorting takes time m * log(m), and the binary searching that np.searchsorted does takes time n * log(m), for a total of (n + m) * log(m). You therefore want to sort the shortest of the two arrays:

def find_rows(a, b):
    dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))

    a_view = np.ascontiguousarray(a).view(dt).ravel()
    b_view = np.ascontiguousarray(b).view(dt).ravel()

    sort_b = np.argsort(b_view)
    where_in_b = np.searchsorted(b_view, a_view,
                                 sorter=sort_b)
    where_in_b = np.take(sort_b, where_in_b)
    which_in_a = np.take(b_view, where_in_b) == a_view
    where_in_b = where_in_b[which_in_a]
    which_in_a = np.nonzero(which_in_a)[0]
    return np.column_stack((which_in_a, where_in_b))

With a and b your two sample arrays:

In [14]: find_rows(a, b)
Out[14]: 
array([[0, 4],
       [2, 1],
       [3, 2],
       [4, 3]], dtype=int64)

In [15]: %timeit find_rows(a, b)
10000 loops, best of 3: 29.7 us per loop

On my system the dictionary approach clocks faster at about 22 us for your test data, but with arrays of 1000x4, this numpy approach is about 6x faster than the pure Python one (483 us vs 2.54 ms).