Get 2d np.array row indices depending on other 2d np.array-CodePudding

I want to fetch all row indicies of an 2d np.array on the basis of another 2d np.array.

Input:
all_elements = np.array([[1, 1], [2, 2], [3, 3]])
elements = np.array([[1, 1],[2, 2]])

Something like:

idx = np.row_idx(elements, all_elements, axis=0)
output:
[0, 1]

I have been trying to do this with some version of np.where(np.isin(.....)) but i can't get it to work properly. Any1 have any tips?

CodePudding user response：

IIUC, you might want to use:

np.where((all_elements==elements[:,None]).all(2).any(0))[0]

output: array([0, 1])

explanation:

# compare all elements using broadcasting
(all_elements==elements[:,None])

array([[[ True,  True],
        [False, False],
        [False, False]],

       [[False, False],
        [ True,  True],
        [False, False]]])

# all True on the last dimension
(all_elements==elements[:,None]).all(2)

array([[ True, False, False],
       [False,  True, False]])

# any match per first dimension
(all_elements==elements[:,None]).all(2).any(0)

array([ True,  True, False])