Get rows of a first array matching rows of a second one

Question

Suppose I have an array A of shape (M, K) and another B of shape (N, K).

The rows of B are all the possible patterns that can be encountered (each pattern is thus a 1D array of size K).

I thus would like to get an array C of shape (M,) where C[i] is the indice of the pattern (in B) of row i in A.

I am currently doing this in a loop (i.e. looping over all the possible patterns) but I would end up using vectorization.

Here is an example:

A = np.array([[0, 1], [0, 1], [1, 0]])
B = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])

I am expecting:

C = np.array([1, 1, 2])

Answer 1

Based on this solution, here's a vectorized solution using np.searchsorted -

dims = B.max(0)+1
A1D = np.ravel_multi_index(A.T,dims)
B1D = np.ravel_multi_index(B.T,dims)
sidx = B1D.argsort()
out = sidx[np.searchsorted(B1D,A1D,sorter=sidx)]

Sample run -

In [43]: A
Out[43]: 
array([[72, 89, 75],
       [72, 89, 75],
       [93, 38, 61],
       [47, 67, 50],
       [47, 67, 50],
       [93, 38, 61],
       [72, 89, 75]])

In [44]: B
Out[44]: 
array([[47, 67, 50],
       [93, 38, 61],
       [41, 55, 27],
       [72, 89, 75]])

In [45]: out
Out[45]: array([3, 3, 1, 0, 0, 1, 3])