python - Numpy 获取二维数组中 arange 的索引

Question

考虑以下 numpy 数组:

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

我希望能够从数组中获取 np.arange(4) 的索引。即在第 0 行中获取索引 0，在第 1 行中获取索引 1，依此类推。

即对于这个具体的例子:

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

在 numpy 中是否有比循环遍历每一行并获取索引更有效的方法:

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

Answer 1

试试这个

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

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