Как получить индекс конкретной строки, используя значения столбцов в numpy?

Question

Итак, если у меня есть следующий массив массива:

>>> arr
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

Теперь, если я хочу получить первую строку, я бы сделал что-то вроде этого:

>>> arr[0]
array([0, 1, 2, 3, 4])

Однако, когда Я использую np.where, чтобы найти определенную строку, например:

>>> np.where(arr == [0,1,2,3,4])

Я получаю этот вывод!

(array([0, 0, 0, 0, 0, 3, 3, 3, 3, 3], dtype=int64), 
array([0, 1, 2, 3, 4, 0, 1, 2, 3, 4], dtype=int64))

Тем не менее, это не то, что я после. Я хотел бы получить индексы строк вместо этого. Например:

(array([0, 3], dtype=int64)

Есть ли способ достичь этого? Любой совет очень ценится!

Answer 1

Я думаю, вы хотите проверить, равны ли строки заданному массиву. В этом случае вам нужно all:

np.where((arr == [0,1,2,3,4]).all(1))
# (array([0, 3]),)