floflo29 floflo29 - 1 month ago 9
Python Question

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

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

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