Numpy get index of arange in 2d array

Question

Consider the following numpy array:

import numpy as np
arr = np.array([np.random.permutation(4) for _ in range(4)])

array([[0, 1, 2, 3],
       [3, 1, 0, 2],
       [1, 2, 0, 3],
       [0, 2, 3, 1]])

I would like to be able to get the index of np.arange(4) from the array. i.e get index of 0 in row 0, index of 1 in row 1, and so on.

i.e for this specific example:

array([0, 1, 1, 2])

Is there a more efficient way to do that in numpy than just looping over each row and getting the index:

alist = []
for ridx in range(arr.shape[0]):
    alist.append(arr[ridx].tolist().index(ridx))
ans = np.array(alist)

Answer 1

Try this

np.nonzero(arr == np.arange(arr.shape[0])[:,None])[1]

Out[15]: array([0, 1, 1, 2], dtype=int64)