Numpy Where function for multidimensional arrays

Question

I have two arrays of the shape (N1, 2) and (N2, 2):

import numpy as np
arr1 = np.array([[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]])
arr2 = np.array([[0, 2], [0, 3], [1, 2], [1, 2]])

I know that every element in arr1 is unique. The elements in arr2 do not have to be unique, and each element in arr2 appears exactly once in arr1.

What I want to do, I want to obtain an array of indices of shape(N2,1) that just gives me the position of the elements in arr2 in arr1. So in this case I want the output to be

idxs = np.array([1, 2, 3, 3])

So in other words, I am looking for something like np.where that works with multidimensional arrays, but the np.where does not have an axis keyword.

I could go for something like

idxs = []
for element in arr2: 
    for i, other_element in enumerate(arr1):
        if (element == other_element).all():
            idxs.append(i)
            break   

But this is ugly and I bet there's a neat numpy solution.

Answer 1

You can use broadcasting to check for membership for each row in arr2, all to check in which rows all elements match and argmax to get the corresponding indices:

((arr2[:,None] == arr1).all(2)).argmax(1)
# array([1, 2, 3, 3], dtype=int64)